From 2c6d4d4dee169db6a95efc786c273ed027ec25cc Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Thu, 16 Jul 2026 16:01:48 +0530 Subject: [PATCH 1/8] natten: port NATTEN v0.21.6 to kernel-builder Vendors SHI-Labs/NATTEN v0.21.6 (MIT) as a kernel-builder edition-5 CUDA kernel targeting kernels-community/natten: - csrc/ vendored with committed autogen output (144 TUs, upstream's 'default' policy); scripts/regen.sh + generate_build_toml.py make upstream bumps a one-command regen. - pybind11 bindings replaced with TORCH_LIBRARY registrations (torch-ext/torch_binding.cpp, 37 out-variant ops, limited-API safe); torch/extension.h stripped from static and autogen sources. - Full natten Python package (minus profiler) vendored with relative imports; _libnatten rewritten on top of the generated _ops module with fake (meta) registrations for torch.compile. - Kernel sections split by arch: core (reference + CUTLASS 2.X FNA/FMHA, SM70-SM120), hopper (9.0a), blackwell (10.0a, cuda-minver 12.8 so cu126 variants still build; natten.h derives the backend gates from CUDART_VERSION so static dispatchers stub out consistently). - Upstream test suite vendored + fast kernels_ci parity subset. Co-Authored-By: Claude Fable 5 --- natten/CARD.md | 112 + natten/LICENSE | 21 + natten/benchmarks/benchmark.py | 73 + natten/build.toml | 459 + .../cuda/blackwell_fmha/dispatch_head_dim.h | 93 + .../cuda/blackwell_fmha/dispatch_tile_size.h | 209 + .../cuda/blackwell_fmha/interface.h | 48 + .../cuda/blackwell_fmha/kernels.h | 628 + .../blackwell_fmha_bwd/dispatch_head_dim.h | 61 + .../blackwell_fmha_bwd/dispatch_tile_size.h | 95 + .../cuda/blackwell_fmha_bwd/interface.h | 42 + .../cuda/blackwell_fmha_bwd/kernels.h | 208 + .../cuda/blackwell_fna/dispatch_cm.h | 785 + .../cuda/blackwell_fna/dispatch_dtype.h | 86 + .../cuda/blackwell_fna/dispatch_head_dim.h | 221 + .../cuda/blackwell_fna/dispatch_tile_shape.h | 8309 +++ .../cuda/blackwell_fna/interface.h | 45 + .../cuda/blackwell_fna/kernels.h | 44428 ++++++++++++++++ .../cuda/blackwell_fna_bwd/dispatch_cm.h | 407 + .../cuda/blackwell_fna_bwd/dispatch_dtype.h | 68 + .../blackwell_fna_bwd/dispatch_head_dim.h | 125 + .../blackwell_fna_bwd/dispatch_tile_shape.h | 2777 + .../cuda/blackwell_fna_bwd/interface.h | 45 + .../cuda/blackwell_fna_bwd/kernels.h | 15688 ++++++ .../natten_autogen/cuda/fmha/dispatch_dtype.h | 129 + .../natten_autogen/cuda/fmha/interface.h | 57 + .../natten_autogen/cuda/fmha/kernels.h | 1226 + .../natten_autogen/cuda/fna/dispatch_cm.h | 1154 + .../natten_autogen/cuda/fna/dispatch_device.h | 134 + .../natten_autogen/cuda/fna/dispatch_dtype.h | 350 + .../natten_autogen/cuda/fna/interface.h | 52 + .../include/natten_autogen/cuda/fna/kernels.h | 16931 ++++++ .../cuda/hopper_fmha/dispatch_head_dim.h | 68 + .../cuda/hopper_fmha/dispatch_tile_size.h | 136 + .../cuda/hopper_fmha/interface.h | 43 + .../natten_autogen/cuda/hopper_fmha/kernels.h | 269 + .../cuda/hopper_fmha_bwd/dispatch_head_dim.h | 61 + .../cuda/hopper_fmha_bwd/dispatch_tile_size.h | 111 + .../cuda/hopper_fmha_bwd/interface.h | 42 + .../cuda/hopper_fmha_bwd/kernels.h | 308 + .../cuda/hopper_fna/dispatch_cm.h | 534 + .../cuda/hopper_fna/dispatch_dtype.h | 69 + .../cuda/hopper_fna/dispatch_head_dim.h | 144 + .../cuda/hopper_fna/dispatch_tile_shape.h | 2034 + .../cuda/hopper_fna/interface.h | 46 + .../natten_autogen/cuda/hopper_fna/kernels.h | 5397 ++ .../cuda/hopper_fna_bwd/dispatch_cm.h | 407 + .../cuda/hopper_fna_bwd/dispatch_dtype.h | 68 + .../cuda/hopper_fna_bwd/dispatch_head_dim.h | 125 + .../cuda/hopper_fna_bwd/dispatch_tile_shape.h | 1785 + .../cuda/hopper_fna_bwd/interface.h | 45 + .../cuda/hopper_fna_bwd/kernels.h | 7620 +++ .../cuda/reference/dispatch_cm.h | 405 + .../cuda/reference/dispatch_dtype.h | 123 + .../natten_autogen/cuda/reference/interface.h | 59 + .../natten_autogen/cuda/reference/kernels.h | 1874 + .../src/cuda/blackwell_fmha/source_0.cu | 634 + .../src/cuda/blackwell_fmha/source_1.cu | 634 + .../src/cuda/blackwell_fmha/source_2.cu | 634 + .../src/cuda/blackwell_fmha/source_3.cu | 634 + .../src/cuda/blackwell_fmha_bwd/source_0.cu | 288 + .../src/cuda/blackwell_fmha_bwd/source_1.cu | 288 + .../src/cuda/blackwell_fmha_bwd/source_2.cu | 155 + .../src/cuda/blackwell_fmha_bwd/source_3.cu | 155 + .../src/cuda/blackwell_fna/source_0.cu | 4118 ++ .../src/cuda/blackwell_fna/source_1.cu | 4118 ++ .../src/cuda/blackwell_fna/source_10.cu | 4118 ++ .../src/cuda/blackwell_fna/source_11.cu | 4118 ++ .../src/cuda/blackwell_fna/source_12.cu | 4054 ++ .../src/cuda/blackwell_fna/source_13.cu | 4054 ++ .../src/cuda/blackwell_fna/source_14.cu | 4054 ++ .../src/cuda/blackwell_fna/source_15.cu | 4054 ++ .../src/cuda/blackwell_fna/source_16.cu | 4054 ++ .../src/cuda/blackwell_fna/source_17.cu | 4054 ++ .../src/cuda/blackwell_fna/source_18.cu | 4054 ++ .../src/cuda/blackwell_fna/source_19.cu | 4054 ++ .../src/cuda/blackwell_fna/source_2.cu | 4118 ++ .../src/cuda/blackwell_fna/source_20.cu | 4054 ++ .../src/cuda/blackwell_fna/source_21.cu | 4054 ++ .../src/cuda/blackwell_fna/source_22.cu | 4054 ++ .../src/cuda/blackwell_fna/source_23.cu | 4054 ++ .../src/cuda/blackwell_fna/source_24.cu | 4054 ++ .../src/cuda/blackwell_fna/source_25.cu | 4054 ++ .../src/cuda/blackwell_fna/source_26.cu | 4054 ++ .../src/cuda/blackwell_fna/source_27.cu | 4054 ++ .../src/cuda/blackwell_fna/source_3.cu | 4118 ++ .../src/cuda/blackwell_fna/source_4.cu | 4118 ++ .../src/cuda/blackwell_fna/source_5.cu | 4118 ++ .../src/cuda/blackwell_fna/source_6.cu | 4118 ++ .../src/cuda/blackwell_fna/source_7.cu | 4118 ++ .../src/cuda/blackwell_fna/source_8.cu | 4118 ++ .../src/cuda/blackwell_fna/source_9.cu | 4118 ++ .../src/cuda/blackwell_fna_bwd/source_0.cu | 2830 + .../src/cuda/blackwell_fna_bwd/source_1.cu | 2830 + .../src/cuda/blackwell_fna_bwd/source_10.cu | 2758 + .../src/cuda/blackwell_fna_bwd/source_11.cu | 2758 + .../src/cuda/blackwell_fna_bwd/source_12.cu | 2758 + .../src/cuda/blackwell_fna_bwd/source_13.cu | 2758 + .../src/cuda/blackwell_fna_bwd/source_2.cu | 2830 + .../src/cuda/blackwell_fna_bwd/source_3.cu | 2830 + .../src/cuda/blackwell_fna_bwd/source_4.cu | 2830 + .../src/cuda/blackwell_fna_bwd/source_5.cu | 2830 + .../src/cuda/blackwell_fna_bwd/source_6.cu | 2830 + .../src/cuda/blackwell_fna_bwd/source_7.cu | 2830 + .../src/cuda/blackwell_fna_bwd/source_8.cu | 2758 + .../src/cuda/blackwell_fna_bwd/source_9.cu | 2758 + natten/csrc/autogen/src/cuda/fmha/source_0.cu | 780 + natten/csrc/autogen/src/cuda/fmha/source_1.cu | 780 + natten/csrc/autogen/src/cuda/fmha/source_2.cu | 780 + natten/csrc/autogen/src/cuda/fmha/source_3.cu | 276 + natten/csrc/autogen/src/cuda/fmha/source_4.cu | 318 + natten/csrc/autogen/src/cuda/fmha/source_5.cu | 402 + natten/csrc/autogen/src/cuda/fna/source_0.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_1.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_10.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_11.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_12.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_13.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_14.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_15.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_16.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_17.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_18.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_19.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_2.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_20.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_21.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_22.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_23.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_24.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_25.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_26.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_27.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_28.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_29.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_3.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_30.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_31.cu | 701 + natten/csrc/autogen/src/cuda/fna/source_32.cu | 364 + natten/csrc/autogen/src/cuda/fna/source_33.cu | 448 + natten/csrc/autogen/src/cuda/fna/source_34.cu | 448 + natten/csrc/autogen/src/cuda/fna/source_35.cu | 532 + natten/csrc/autogen/src/cuda/fna/source_36.cu | 364 + natten/csrc/autogen/src/cuda/fna/source_37.cu | 364 + natten/csrc/autogen/src/cuda/fna/source_38.cu | 364 + natten/csrc/autogen/src/cuda/fna/source_39.cu | 532 + natten/csrc/autogen/src/cuda/fna/source_4.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_40.cu | 364 + natten/csrc/autogen/src/cuda/fna/source_41.cu | 364 + natten/csrc/autogen/src/cuda/fna/source_42.cu | 532 + natten/csrc/autogen/src/cuda/fna/source_43.cu | 532 + natten/csrc/autogen/src/cuda/fna/source_44.cu | 532 + natten/csrc/autogen/src/cuda/fna/source_45.cu | 364 + natten/csrc/autogen/src/cuda/fna/source_46.cu | 364 + natten/csrc/autogen/src/cuda/fna/source_47.cu | 364 + natten/csrc/autogen/src/cuda/fna/source_48.cu | 364 + natten/csrc/autogen/src/cuda/fna/source_49.cu | 364 + natten/csrc/autogen/src/cuda/fna/source_5.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_50.cu | 364 + natten/csrc/autogen/src/cuda/fna/source_51.cu | 532 + natten/csrc/autogen/src/cuda/fna/source_52.cu | 532 + natten/csrc/autogen/src/cuda/fna/source_53.cu | 364 + natten/csrc/autogen/src/cuda/fna/source_54.cu | 364 + natten/csrc/autogen/src/cuda/fna/source_55.cu | 364 + natten/csrc/autogen/src/cuda/fna/source_56.cu | 364 + natten/csrc/autogen/src/cuda/fna/source_57.cu | 532 + natten/csrc/autogen/src/cuda/fna/source_58.cu | 532 + natten/csrc/autogen/src/cuda/fna/source_59.cu | 532 + natten/csrc/autogen/src/cuda/fna/source_6.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_60.cu | 402 + natten/csrc/autogen/src/cuda/fna/source_61.cu | 402 + natten/csrc/autogen/src/cuda/fna/source_62.cu | 402 + natten/csrc/autogen/src/cuda/fna/source_63.cu | 402 + natten/csrc/autogen/src/cuda/fna/source_7.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_8.cu | 1036 + natten/csrc/autogen/src/cuda/fna/source_9.cu | 1036 + .../autogen/src/cuda/hopper_fmha/source_0.cu | 222 + .../autogen/src/cuda/hopper_fmha/source_1.cu | 222 + .../autogen/src/cuda/hopper_fmha/source_2.cu | 222 + .../autogen/src/cuda/hopper_fmha/source_3.cu | 222 + .../autogen/src/cuda/hopper_fmha/source_4.cu | 222 + .../src/cuda/hopper_fmha_bwd/source_0.cu | 234 + .../src/cuda/hopper_fmha_bwd/source_1.cu | 234 + .../src/cuda/hopper_fmha_bwd/source_2.cu | 234 + .../src/cuda/hopper_fmha_bwd/source_3.cu | 234 + .../src/cuda/hopper_fmha_bwd/source_4.cu | 234 + .../autogen/src/cuda/hopper_fna/source_0.cu | 1882 + .../autogen/src/cuda/hopper_fna/source_1.cu | 1882 + .../autogen/src/cuda/hopper_fna/source_2.cu | 1882 + .../autogen/src/cuda/hopper_fna/source_3.cu | 1882 + .../autogen/src/cuda/hopper_fna/source_4.cu | 1822 + .../autogen/src/cuda/hopper_fna/source_5.cu | 1822 + .../autogen/src/cuda/hopper_fna/source_6.cu | 1822 + .../autogen/src/cuda/hopper_fna/source_7.cu | 1822 + .../src/cuda/hopper_fna_bwd/source_0.cu | 4986 ++ .../src/cuda/hopper_fna_bwd/source_1.cu | 4986 ++ .../src/cuda/hopper_fna_bwd/source_2.cu | 4986 ++ .../src/cuda/hopper_fna_bwd/source_3.cu | 4986 ++ .../autogen/src/cuda/reference/source_0.cu | 2181 + .../autogen/src/cuda/reference/source_1.cu | 2229 + natten/csrc/include/natten/blackwell_fmha.h | 74 + natten/csrc/include/natten/blackwell_fna.h | 155 + natten/csrc/include/natten/compute_delta.h | 40 + .../cuda/fmha/epilogue/epilogue_pipelined.h | 638 + .../fmha/epilogue/epilogue_rescale_output.h | 267 + .../epilogue_thread_apply_logsumexp.h | 183 + .../natten/cuda/fmha/fmha_backward.cuh | 239 + .../include/natten/cuda/fmha/fmha_forward.cuh | 191 + .../natten/cuda/fmha/gemm/custom_mma.h | 131 + .../natten/cuda/fmha/gemm/custom_mma_base.h | 189 + .../cuda/fmha/gemm/custom_mma_multistage.h | 765 + .../cuda/fmha/gemm/custom_mma_pipelined.h | 408 + .../natten/cuda/fmha/gemm/find_default_mma.h | 167 + .../fmha/gemm/mma_accum_lambda_iterator.h | 354 + .../natten/cuda/fmha/gemm/mma_from_smem.h | 1916 + .../natten/cuda/fmha/gemm_kernel_utils.h | 240 + .../default_warp_iterator_from_smem.h | 149 + .../epilogue_predicated_tile_iterator.h | 760 + .../cuda/fmha/iterators/make_residual_last.h | 73 + ...cated_tile_access_iterator_residual_last.h | 2122 + .../predicated_tile_iterator_residual_last.h | 2127 + .../fmha/iterators/transpose_warp_iterator.h | 60 + .../fmha/iterators/warp_iterator_from_smem.h | 290 + .../natten/cuda/fmha/kernel_backward.h | 2282 + .../include/natten/cuda/fmha/kernel_forward.h | 1075 + .../cuda/fmha/transform/tile_smem_loader.h | 66 + .../fmha_blackwell/collective/fmha_common.hpp | 180 + .../fmha_blackwell/collective/fmha_fusion.hpp | 379 + ..._fmha_fwd_epilogue_tma_warpspecialized.hpp | 240 + ..._fmha_fwd_mainloop_tma_warpspecialized.hpp | 1387 + .../sm100_fmha_load_tma_warpspecialized.hpp | 346 + .../cuda/fmha_blackwell/common/pow_2.hpp | 92 + .../fmha_blackwell/device/fmha_bwd_sm100.hpp | 446 + .../cuda/fmha_blackwell/device/fmha_sm100.hpp | 286 + .../cuda/fmha_blackwell/fmha_backward.cuh | 255 + .../cuda/fmha_blackwell/fmha_forward.cuh | 276 + .../kernel/fmha_causal_tile_scheduler.hpp | 218 + .../kernel/fmha_kernel_bwd_convert.hpp | 218 + .../kernel/fmha_kernel_bwd_sum_OdO.hpp | 194 + .../kernel/fmha_tile_scheduler.hpp | 176 + ...00_fmha_bwd_kernel_tma_warpspecialized.hpp | 2334 + ...00_fmha_fwd_kernel_tma_warpspecialized.hpp | 757 + ...mha_collective_bwd_tma_warpspecialized.hpp | 1235 + .../collective/fmha_collective_load.hpp | 223 + .../collective/fmha_collective_softmax.hpp | 378 + .../collective/fmha_collective_tma.hpp | 715 + .../fmha_collective_tma_warpspecialized.hpp | 775 + .../fmha_hopper/collective/fmha_common.hpp | 317 + .../fmha_hopper/collective/fmha_epilogue.hpp | 239 + .../collective/fmha_epilogue_bwd.hpp | 259 + .../fmha_hopper/collective/fmha_fusion.hpp | 281 + .../fmha_hopper/collective/fmha_varlen.hpp | 109 + .../cuda/fmha_hopper/device/fmha_bwd_sm90.hpp | 400 + .../cuda/fmha_hopper/device/fmha_sm90.hpp | 288 + .../natten/cuda/fmha_hopper/fmha_backward.cuh | 224 + .../natten/cuda/fmha_hopper/fmha_forward.cuh | 269 + .../kernel/fmha_kernel_builder.hpp | 195 + .../kernel/fmha_kernel_bwd_convert.hpp | 210 + .../kernel/fmha_kernel_bwd_sum_OdO.hpp | 192 + .../fmha_hopper/kernel/fmha_kernel_tma.hpp | 255 + .../fmha_kernel_tma_warpspecialized.hpp | 491 + .../cuda/fmha_hopper/kernel/fmha_options.hpp | 83 + .../kernel/fmha_tile_scheduler.hpp | 227 + .../cuda/fna/epilogue/epilogue_pipelined.h | 638 + .../fna/epilogue/epilogue_rescale_output.h | 267 + .../epilogue_thread_apply_logsumexp.h | 184 + .../fna/epilogue/predicated_tile_iterator.h | 628 + .../predicated_tile_iterator_params.h | 152 + .../include/natten/cuda/fna/fna_backward.cuh | 219 + .../include/natten/cuda/fna/fna_forward.cuh | 192 + .../include/natten/cuda/fna/gemm/custom_mma.h | 271 + .../natten/cuda/fna/gemm/custom_mma_base.h | 189 + .../cuda/fna/gemm/custom_mma_multistage.h | 765 + .../cuda/fna/gemm/custom_mma_pipelined.h | 408 + .../natten/cuda/fna/gemm/find_default_mma.h | 272 + .../cuda/fna/gemm/mma_accum_lambda_iterator.h | 393 + .../natten/cuda/fna/gemm/mma_from_smem.h | 1939 + .../cuda/fna/gemm/replace_mma_iterators.h | 243 + .../natten/cuda/fna/gemm_kernel_utils.h | 186 + .../default_warp_iterator_from_smem.h | 149 + .../epilogue_predicated_tile_iterator.h | 562 + .../cuda/fna/iterators/make_residual_last.h | 108 + .../predicated_tile_access_iterator.h | 1662 + ...cated_tile_access_iterator_residual_last.h | 976 + .../fna/iterators/predicated_tile_iterator.h | 831 + .../predicated_tile_iterator_residual_last.h | 617 + .../fna/iterators/transpose_warp_iterator.h | 60 + .../fna/iterators/warp_iterator_from_smem.h | 290 + .../include/natten/cuda/fna/kernel_backward.h | 2763 + .../include/natten/cuda/fna/kernel_forward.h | 1138 + .../csrc/include/natten/cuda/fna/na_utils.cuh | 1725 + .../fna_blackwell/collective/fna_common.hpp | 180 + .../fna_blackwell/collective/fna_fusion.hpp | 609 + .../collective/fna_fusion_bwd.hpp | 450 + ...0_fna_fwd_epilogue_tma_warpspecialized.hpp | 215 + ...0_fna_fwd_mainloop_tma_warpspecialized.hpp | 1454 + .../sm100_fna_load_tma_warpspecialized.hpp | 393 + .../cuda/fna_blackwell/common/pow_2.hpp | 92 + .../fna_blackwell/device/fna_bwd_sm100.hpp | 433 + .../cuda/fna_blackwell/device/fna_sm100.hpp | 286 + .../cuda/fna_blackwell/fna_backward.cuh | 249 + .../natten/cuda/fna_blackwell/fna_forward.cuh | 270 + ...100_fna_bwd_kernel_tma_warpspecialized.hpp | 2389 + ...100_fna_fwd_kernel_tma_warpspecialized.hpp | 713 + ...fna_collective_bwd_tma_warpspecialized.hpp | 1422 + .../collective/fna_collective_load.hpp | 232 + .../collective/fna_collective_softmax.hpp | 479 + .../collective/fna_collective_tma.hpp | 828 + .../fna_collective_tma_warpspecialized.hpp | 906 + .../cuda/fna_hopper/collective/fna_common.hpp | 360 + .../cuda/fna_hopper/collective/fna_fusion.hpp | 543 + .../fna_hopper/collective/fna_fusion_bwd.hpp | 450 + .../cuda/fna_hopper/device/fna_bwd_sm90.hpp | 426 + .../cuda/fna_hopper/device/fna_sm90.hpp | 287 + .../natten/cuda/fna_hopper/fna_backward.cuh | 200 + .../natten/cuda/fna_hopper/fna_forward.cuh | 252 + .../fna_hopper/kernel/fna_kernel_builder.hpp | 232 + .../include/natten/cuda/hopper_fmha_fna.h | 71 + .../natten/cuda/reduction/compute_delta.cuh | 99 + .../reduction/fmha_kernel_bwd_sum_OdO.hpp | 160 + .../cuda/reference/fna_reference_backward.hpp | 773 + .../cuda/reference/fna_reference_forward.hpp | 329 + .../include/natten/cuda/reference/mask.hpp | 353 + .../include/natten/cuda/reference/utils.hpp | 65 + .../include/natten/cuda/tokperm/layouts.hpp | 150 + .../cuda/tokperm/token_permute_kernel.cuh | 355 + .../include/natten/cuda/tokperm/tokperm.hpp | 262 + .../natten/cuda/tokperm/utils/permute.cuh | 57 + .../natten/cuda/tokperm/utils/stride.cuh | 75 + .../natten/cuda/tokperm/utils/tuple.cuh | 40 + natten/csrc/include/natten/cuda/utils/cuda.h | 37 + .../include/natten/cuda/utils/cutlass.cuh | 112 + .../cuda/utils/generic_cutlass_device.hpp | 288 + natten/csrc/include/natten/fmha.h | 74 + natten/csrc/include/natten/fna.h | 140 + natten/csrc/include/natten/helpers.h | 461 + natten/csrc/include/natten/hopper_fmha.h | 73 + natten/csrc/include/natten/hopper_fna.h | 155 + natten/csrc/include/natten/natten.h | 147 + natten/csrc/include/natten/reference.h | 134 + natten/csrc/include/natten/token_permute.h | 76 + natten/csrc/src/blackwell_fmha.cu | 495 + natten/csrc/src/blackwell_fna_backward.cu | 406 + natten/csrc/src/blackwell_fna_forward.cu | 383 + natten/csrc/src/compute_delta.cu | 123 + natten/csrc/src/fmha.cu | 424 + natten/csrc/src/fna_backward.cu | 345 + natten/csrc/src/fna_forward.cu | 238 + natten/csrc/src/hopper_fmha.cu | 409 + natten/csrc/src/hopper_fna_backward.cu | 358 + natten/csrc/src/hopper_fna_forward.cu | 327 + natten/csrc/src/reference_backward.cu | 335 + natten/csrc/src/reference_forward.cu | 286 + natten/csrc/src/token_permute.cu | 268 + natten/csrc/src/token_unpermute.cu | 268 + natten/flake.lock | 117 + natten/flake.nix | 17 + natten/scripts/autogen_blackwell_fmha.py | 607 + natten/scripts/autogen_blackwell_fmha_bwd.py | 623 + natten/scripts/autogen_blackwell_fna.py | 793 + natten/scripts/autogen_blackwell_fna_bwd.py | 790 + natten/scripts/autogen_fmha.py | 676 + natten/scripts/autogen_fna.py | 906 + natten/scripts/autogen_hopper_fmha.py | 669 + natten/scripts/autogen_hopper_fmha_bwd.py | 620 + natten/scripts/autogen_hopper_fna.py | 885 + natten/scripts/autogen_hopper_fna_bwd.py | 838 + natten/scripts/autogen_reference_fna.py | 621 + natten/scripts/generate_build_toml.py | 145 + natten/scripts/regen.sh | 38 + natten/tests/__init__.py | 22 + natten/tests/conftest.py | 10 + natten/tests/test_attn_merge.py | 286 + natten/tests/test_blackwell_fna.py | 556 + natten/tests/test_compute_delta.py | 125 + natten/tests/test_fmha.py | 1449 + natten/tests/test_fmha_varlen.py | 980 + natten/tests/test_fna.py | 791 + natten/tests/test_hopper_fna.py | 549 + natten/tests/test_kernels_ci.py | 110 + natten/tests/test_token_permute.py | 332 + natten/tests/test_torch_compile.py | 575 + natten/tests/utils.py | 319 + natten/torch-ext/natten/__init__.py | 106 + natten/torch-ext/natten/_environment.py | 59 + .../torch-ext/natten/_libnatten/__init__.py | 109 + .../natten/_libnatten/torch_wrappers.py | 1006 + natten/torch-ext/natten/attn_merge.py | 292 + natten/torch-ext/natten/backends/__init__.py | 254 + .../natten/backends/blackwell_fmha.py | 254 + .../natten/backends/blackwell_fna.py | 500 + .../natten/backends/configs/__init__.py | 584 + .../natten/backends/configs/checks.py | 750 + .../backends/configs/cutlass/__init__.py | 430 + .../configs/cutlass/backward_knobs.py | 228 + .../configs/cutlass/fna_backward_128x128.py | 304 + .../configs/cutlass/fna_backward_128x64.py | 223 + .../configs/cutlass/fna_backward_64x64.py | 168 + .../configs/cutlass/fna_forward_32x128.py | 90 + .../configs/cutlass/fna_forward_64x128.py | 82 + .../configs/cutlass/fna_forward_64x64.py | 63 + .../configs/cutlass_blackwell/__init__.py | 391 + .../configs/cutlass_hopper/__init__.py | 522 + .../natten/backends/configs/flex/__init__.py | 210 + natten/torch-ext/natten/backends/flex.py | 799 + natten/torch-ext/natten/backends/fmha.py | 283 + natten/torch-ext/natten/backends/fna.py | 417 + .../torch-ext/natten/backends/hopper_fmha.py | 261 + .../torch-ext/natten/backends/hopper_fna.py | 512 + natten/torch-ext/natten/backends/reference.py | 343 + natten/torch-ext/natten/context.py | 231 + natten/torch-ext/natten/functional.py | 1151 + natten/torch-ext/natten/modules.py | 449 + .../natten/token_permute/__init__.py | 32 + .../natten/token_permute/cutlass_impl.py | 286 + .../natten/token_permute/frontend.py | 137 + .../natten/token_permute/torch_impl.py | 368 + natten/torch-ext/natten/types.py | 85 + natten/torch-ext/natten/utils/__init__.py | 22 + natten/torch-ext/natten/utils/checks.py | 726 + natten/torch-ext/natten/utils/device.py | 50 + natten/torch-ext/natten/utils/dtype.py | 36 + natten/torch-ext/natten/utils/environment.py | 79 + natten/torch-ext/natten/utils/log.py | 134 + natten/torch-ext/natten/utils/tensor.py | 113 + natten/torch-ext/natten/utils/testing.py | 149 + natten/torch-ext/natten/utils/tuples.py | 51 + natten/torch-ext/natten/utils/varlen.py | 135 + natten/torch-ext/natten/version.py | 24 + natten/torch-ext/torch_binding.cpp | 871 + 430 files changed, 465236 insertions(+) create mode 100644 natten/CARD.md create mode 100644 natten/LICENSE create mode 100644 natten/benchmarks/benchmark.py create mode 100644 natten/build.toml create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/dispatch_head_dim.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/dispatch_tile_size.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/interface.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/kernels.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/dispatch_head_dim.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/dispatch_tile_size.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/interface.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/kernels.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_cm.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_dtype.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_head_dim.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_tile_shape.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna/interface.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna/kernels.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_cm.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_dtype.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_head_dim.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_tile_shape.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/interface.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/kernels.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/fmha/dispatch_dtype.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/fmha/interface.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/fmha/kernels.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/fna/dispatch_cm.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/fna/dispatch_device.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/fna/dispatch_dtype.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/fna/interface.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/fna/kernels.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha/dispatch_head_dim.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha/dispatch_tile_size.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha/interface.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha/kernels.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/dispatch_head_dim.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/dispatch_tile_size.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/interface.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/kernels.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_cm.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_dtype.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_head_dim.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_tile_shape.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna/interface.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna/kernels.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_cm.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_dtype.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_head_dim.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_tile_shape.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/interface.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/kernels.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/reference/dispatch_cm.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/reference/dispatch_dtype.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/reference/interface.h create mode 100644 natten/csrc/autogen/include/natten_autogen/cuda/reference/kernels.h create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fmha/source_0.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fmha/source_1.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fmha/source_2.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fmha/source_3.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_0.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_1.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_2.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_3.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna/source_0.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna/source_1.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna/source_10.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna/source_11.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna/source_12.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna/source_13.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna/source_14.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna/source_15.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna/source_16.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna/source_17.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna/source_18.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna/source_19.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna/source_2.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna/source_20.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna/source_21.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna/source_22.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna/source_23.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna/source_24.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna/source_25.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna/source_26.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna/source_27.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna/source_3.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna/source_4.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna/source_5.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna/source_6.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna/source_7.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna/source_8.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna/source_9.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_0.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_1.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_10.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_11.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_12.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_13.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_2.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_3.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_4.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_5.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_6.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_7.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_8.cu create mode 100644 natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_9.cu create mode 100644 natten/csrc/autogen/src/cuda/fmha/source_0.cu create mode 100644 natten/csrc/autogen/src/cuda/fmha/source_1.cu create mode 100644 natten/csrc/autogen/src/cuda/fmha/source_2.cu create mode 100644 natten/csrc/autogen/src/cuda/fmha/source_3.cu create mode 100644 natten/csrc/autogen/src/cuda/fmha/source_4.cu create mode 100644 natten/csrc/autogen/src/cuda/fmha/source_5.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_0.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_1.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_10.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_11.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_12.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_13.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_14.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_15.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_16.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_17.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_18.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_19.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_2.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_20.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_21.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_22.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_23.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_24.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_25.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_26.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_27.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_28.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_29.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_3.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_30.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_31.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_32.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_33.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_34.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_35.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_36.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_37.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_38.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_39.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_4.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_40.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_41.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_42.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_43.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_44.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_45.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_46.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_47.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_48.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_49.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_5.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_50.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_51.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_52.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_53.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_54.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_55.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_56.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_57.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_58.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_59.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_6.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_60.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_61.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_62.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_63.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_7.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_8.cu create mode 100644 natten/csrc/autogen/src/cuda/fna/source_9.cu create mode 100644 natten/csrc/autogen/src/cuda/hopper_fmha/source_0.cu create mode 100644 natten/csrc/autogen/src/cuda/hopper_fmha/source_1.cu create mode 100644 natten/csrc/autogen/src/cuda/hopper_fmha/source_2.cu create mode 100644 natten/csrc/autogen/src/cuda/hopper_fmha/source_3.cu create mode 100644 natten/csrc/autogen/src/cuda/hopper_fmha/source_4.cu create mode 100644 natten/csrc/autogen/src/cuda/hopper_fmha_bwd/source_0.cu create mode 100644 natten/csrc/autogen/src/cuda/hopper_fmha_bwd/source_1.cu create mode 100644 natten/csrc/autogen/src/cuda/hopper_fmha_bwd/source_2.cu create mode 100644 natten/csrc/autogen/src/cuda/hopper_fmha_bwd/source_3.cu create mode 100644 natten/csrc/autogen/src/cuda/hopper_fmha_bwd/source_4.cu create mode 100644 natten/csrc/autogen/src/cuda/hopper_fna/source_0.cu create mode 100644 natten/csrc/autogen/src/cuda/hopper_fna/source_1.cu create mode 100644 natten/csrc/autogen/src/cuda/hopper_fna/source_2.cu create mode 100644 natten/csrc/autogen/src/cuda/hopper_fna/source_3.cu create mode 100644 natten/csrc/autogen/src/cuda/hopper_fna/source_4.cu create mode 100644 natten/csrc/autogen/src/cuda/hopper_fna/source_5.cu create mode 100644 natten/csrc/autogen/src/cuda/hopper_fna/source_6.cu create mode 100644 natten/csrc/autogen/src/cuda/hopper_fna/source_7.cu create mode 100644 natten/csrc/autogen/src/cuda/hopper_fna_bwd/source_0.cu create mode 100644 natten/csrc/autogen/src/cuda/hopper_fna_bwd/source_1.cu create mode 100644 natten/csrc/autogen/src/cuda/hopper_fna_bwd/source_2.cu create mode 100644 natten/csrc/autogen/src/cuda/hopper_fna_bwd/source_3.cu create mode 100644 natten/csrc/autogen/src/cuda/reference/source_0.cu create mode 100644 natten/csrc/autogen/src/cuda/reference/source_1.cu create mode 100644 natten/csrc/include/natten/blackwell_fmha.h create mode 100644 natten/csrc/include/natten/blackwell_fna.h create mode 100644 natten/csrc/include/natten/compute_delta.h create mode 100644 natten/csrc/include/natten/cuda/fmha/epilogue/epilogue_pipelined.h create mode 100644 natten/csrc/include/natten/cuda/fmha/epilogue/epilogue_rescale_output.h create mode 100644 natten/csrc/include/natten/cuda/fmha/epilogue/epilogue_thread_apply_logsumexp.h create mode 100644 natten/csrc/include/natten/cuda/fmha/fmha_backward.cuh create mode 100644 natten/csrc/include/natten/cuda/fmha/fmha_forward.cuh create mode 100644 natten/csrc/include/natten/cuda/fmha/gemm/custom_mma.h create mode 100644 natten/csrc/include/natten/cuda/fmha/gemm/custom_mma_base.h create mode 100644 natten/csrc/include/natten/cuda/fmha/gemm/custom_mma_multistage.h create mode 100644 natten/csrc/include/natten/cuda/fmha/gemm/custom_mma_pipelined.h create mode 100644 natten/csrc/include/natten/cuda/fmha/gemm/find_default_mma.h create mode 100644 natten/csrc/include/natten/cuda/fmha/gemm/mma_accum_lambda_iterator.h create mode 100644 natten/csrc/include/natten/cuda/fmha/gemm/mma_from_smem.h create mode 100644 natten/csrc/include/natten/cuda/fmha/gemm_kernel_utils.h create mode 100644 natten/csrc/include/natten/cuda/fmha/iterators/default_warp_iterator_from_smem.h create mode 100644 natten/csrc/include/natten/cuda/fmha/iterators/epilogue_predicated_tile_iterator.h create mode 100644 natten/csrc/include/natten/cuda/fmha/iterators/make_residual_last.h create mode 100644 natten/csrc/include/natten/cuda/fmha/iterators/predicated_tile_access_iterator_residual_last.h create mode 100644 natten/csrc/include/natten/cuda/fmha/iterators/predicated_tile_iterator_residual_last.h create mode 100644 natten/csrc/include/natten/cuda/fmha/iterators/transpose_warp_iterator.h create mode 100644 natten/csrc/include/natten/cuda/fmha/iterators/warp_iterator_from_smem.h create mode 100644 natten/csrc/include/natten/cuda/fmha/kernel_backward.h create mode 100644 natten/csrc/include/natten/cuda/fmha/kernel_forward.h create mode 100644 natten/csrc/include/natten/cuda/fmha/transform/tile_smem_loader.h create mode 100644 natten/csrc/include/natten/cuda/fmha_blackwell/collective/fmha_common.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_blackwell/collective/fmha_fusion.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_blackwell/collective/sm100_fmha_fwd_epilogue_tma_warpspecialized.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_blackwell/collective/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_blackwell/collective/sm100_fmha_load_tma_warpspecialized.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_blackwell/common/pow_2.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_blackwell/device/fmha_bwd_sm100.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_blackwell/device/fmha_sm100.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_blackwell/fmha_backward.cuh create mode 100644 natten/csrc/include/natten/cuda/fmha_blackwell/fmha_forward.cuh create mode 100644 natten/csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_causal_tile_scheduler.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_kernel_bwd_convert.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_kernel_bwd_sum_OdO.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_tile_scheduler.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_blackwell/kernel/sm100_fmha_bwd_kernel_tma_warpspecialized.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_blackwell/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_bwd_tma_warpspecialized.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_load.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_softmax.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_tma.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_tma_warpspecialized.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_common.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_epilogue.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_epilogue_bwd.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_fusion.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_varlen.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_hopper/device/fmha_bwd_sm90.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_hopper/device/fmha_sm90.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_hopper/fmha_backward.cuh create mode 100644 natten/csrc/include/natten/cuda/fmha_hopper/fmha_forward.cuh create mode 100644 natten/csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_builder.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_bwd_convert.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_bwd_sum_OdO.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_tma.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_tma_warpspecialized.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_hopper/kernel/fmha_options.hpp create mode 100644 natten/csrc/include/natten/cuda/fmha_hopper/kernel/fmha_tile_scheduler.hpp create mode 100644 natten/csrc/include/natten/cuda/fna/epilogue/epilogue_pipelined.h create mode 100644 natten/csrc/include/natten/cuda/fna/epilogue/epilogue_rescale_output.h create mode 100644 natten/csrc/include/natten/cuda/fna/epilogue/epilogue_thread_apply_logsumexp.h create mode 100644 natten/csrc/include/natten/cuda/fna/epilogue/predicated_tile_iterator.h create mode 100644 natten/csrc/include/natten/cuda/fna/epilogue/predicated_tile_iterator_params.h create mode 100644 natten/csrc/include/natten/cuda/fna/fna_backward.cuh create mode 100644 natten/csrc/include/natten/cuda/fna/fna_forward.cuh create mode 100644 natten/csrc/include/natten/cuda/fna/gemm/custom_mma.h create mode 100644 natten/csrc/include/natten/cuda/fna/gemm/custom_mma_base.h create mode 100644 natten/csrc/include/natten/cuda/fna/gemm/custom_mma_multistage.h create mode 100644 natten/csrc/include/natten/cuda/fna/gemm/custom_mma_pipelined.h create mode 100644 natten/csrc/include/natten/cuda/fna/gemm/find_default_mma.h create mode 100644 natten/csrc/include/natten/cuda/fna/gemm/mma_accum_lambda_iterator.h create mode 100644 natten/csrc/include/natten/cuda/fna/gemm/mma_from_smem.h create mode 100644 natten/csrc/include/natten/cuda/fna/gemm/replace_mma_iterators.h create mode 100644 natten/csrc/include/natten/cuda/fna/gemm_kernel_utils.h create mode 100644 natten/csrc/include/natten/cuda/fna/iterators/default_warp_iterator_from_smem.h create mode 100644 natten/csrc/include/natten/cuda/fna/iterators/epilogue_predicated_tile_iterator.h create mode 100644 natten/csrc/include/natten/cuda/fna/iterators/make_residual_last.h create mode 100644 natten/csrc/include/natten/cuda/fna/iterators/predicated_tile_access_iterator.h create mode 100644 natten/csrc/include/natten/cuda/fna/iterators/predicated_tile_access_iterator_residual_last.h create mode 100644 natten/csrc/include/natten/cuda/fna/iterators/predicated_tile_iterator.h create mode 100644 natten/csrc/include/natten/cuda/fna/iterators/predicated_tile_iterator_residual_last.h create mode 100644 natten/csrc/include/natten/cuda/fna/iterators/transpose_warp_iterator.h create mode 100644 natten/csrc/include/natten/cuda/fna/iterators/warp_iterator_from_smem.h create mode 100644 natten/csrc/include/natten/cuda/fna/kernel_backward.h create mode 100644 natten/csrc/include/natten/cuda/fna/kernel_forward.h create mode 100644 natten/csrc/include/natten/cuda/fna/na_utils.cuh create mode 100644 natten/csrc/include/natten/cuda/fna_blackwell/collective/fna_common.hpp create mode 100644 natten/csrc/include/natten/cuda/fna_blackwell/collective/fna_fusion.hpp create mode 100644 natten/csrc/include/natten/cuda/fna_blackwell/collective/fna_fusion_bwd.hpp create mode 100644 natten/csrc/include/natten/cuda/fna_blackwell/collective/sm100_fna_fwd_epilogue_tma_warpspecialized.hpp create mode 100644 natten/csrc/include/natten/cuda/fna_blackwell/collective/sm100_fna_fwd_mainloop_tma_warpspecialized.hpp create mode 100644 natten/csrc/include/natten/cuda/fna_blackwell/collective/sm100_fna_load_tma_warpspecialized.hpp create mode 100644 natten/csrc/include/natten/cuda/fna_blackwell/common/pow_2.hpp create mode 100644 natten/csrc/include/natten/cuda/fna_blackwell/device/fna_bwd_sm100.hpp create mode 100644 natten/csrc/include/natten/cuda/fna_blackwell/device/fna_sm100.hpp create mode 100644 natten/csrc/include/natten/cuda/fna_blackwell/fna_backward.cuh create mode 100644 natten/csrc/include/natten/cuda/fna_blackwell/fna_forward.cuh create mode 100644 natten/csrc/include/natten/cuda/fna_blackwell/kernel/sm100_fna_bwd_kernel_tma_warpspecialized.hpp create mode 100644 natten/csrc/include/natten/cuda/fna_blackwell/kernel/sm100_fna_fwd_kernel_tma_warpspecialized.hpp create mode 100644 natten/csrc/include/natten/cuda/fna_hopper/collective/fna_collective_bwd_tma_warpspecialized.hpp create mode 100644 natten/csrc/include/natten/cuda/fna_hopper/collective/fna_collective_load.hpp create mode 100644 natten/csrc/include/natten/cuda/fna_hopper/collective/fna_collective_softmax.hpp create mode 100644 natten/csrc/include/natten/cuda/fna_hopper/collective/fna_collective_tma.hpp create mode 100644 natten/csrc/include/natten/cuda/fna_hopper/collective/fna_collective_tma_warpspecialized.hpp create mode 100644 natten/csrc/include/natten/cuda/fna_hopper/collective/fna_common.hpp create mode 100644 natten/csrc/include/natten/cuda/fna_hopper/collective/fna_fusion.hpp create mode 100644 natten/csrc/include/natten/cuda/fna_hopper/collective/fna_fusion_bwd.hpp create mode 100644 natten/csrc/include/natten/cuda/fna_hopper/device/fna_bwd_sm90.hpp create mode 100644 natten/csrc/include/natten/cuda/fna_hopper/device/fna_sm90.hpp create mode 100644 natten/csrc/include/natten/cuda/fna_hopper/fna_backward.cuh create mode 100644 natten/csrc/include/natten/cuda/fna_hopper/fna_forward.cuh create mode 100644 natten/csrc/include/natten/cuda/fna_hopper/kernel/fna_kernel_builder.hpp create mode 100644 natten/csrc/include/natten/cuda/hopper_fmha_fna.h create mode 100644 natten/csrc/include/natten/cuda/reduction/compute_delta.cuh create mode 100644 natten/csrc/include/natten/cuda/reduction/fmha_kernel_bwd_sum_OdO.hpp create mode 100644 natten/csrc/include/natten/cuda/reference/fna_reference_backward.hpp create mode 100644 natten/csrc/include/natten/cuda/reference/fna_reference_forward.hpp create mode 100644 natten/csrc/include/natten/cuda/reference/mask.hpp create mode 100644 natten/csrc/include/natten/cuda/reference/utils.hpp create mode 100644 natten/csrc/include/natten/cuda/tokperm/layouts.hpp create mode 100644 natten/csrc/include/natten/cuda/tokperm/token_permute_kernel.cuh create mode 100644 natten/csrc/include/natten/cuda/tokperm/tokperm.hpp create mode 100644 natten/csrc/include/natten/cuda/tokperm/utils/permute.cuh create mode 100644 natten/csrc/include/natten/cuda/tokperm/utils/stride.cuh create mode 100644 natten/csrc/include/natten/cuda/tokperm/utils/tuple.cuh create mode 100644 natten/csrc/include/natten/cuda/utils/cuda.h create mode 100644 natten/csrc/include/natten/cuda/utils/cutlass.cuh create mode 100644 natten/csrc/include/natten/cuda/utils/generic_cutlass_device.hpp create mode 100644 natten/csrc/include/natten/fmha.h create mode 100644 natten/csrc/include/natten/fna.h create mode 100644 natten/csrc/include/natten/helpers.h create mode 100644 natten/csrc/include/natten/hopper_fmha.h create mode 100644 natten/csrc/include/natten/hopper_fna.h create mode 100644 natten/csrc/include/natten/natten.h create mode 100644 natten/csrc/include/natten/reference.h create mode 100644 natten/csrc/include/natten/token_permute.h create mode 100644 natten/csrc/src/blackwell_fmha.cu create mode 100644 natten/csrc/src/blackwell_fna_backward.cu create mode 100644 natten/csrc/src/blackwell_fna_forward.cu create mode 100644 natten/csrc/src/compute_delta.cu create mode 100644 natten/csrc/src/fmha.cu create mode 100644 natten/csrc/src/fna_backward.cu create mode 100644 natten/csrc/src/fna_forward.cu create mode 100644 natten/csrc/src/hopper_fmha.cu create mode 100644 natten/csrc/src/hopper_fna_backward.cu create mode 100644 natten/csrc/src/hopper_fna_forward.cu create mode 100644 natten/csrc/src/reference_backward.cu create mode 100644 natten/csrc/src/reference_forward.cu create mode 100644 natten/csrc/src/token_permute.cu create mode 100644 natten/csrc/src/token_unpermute.cu create mode 100644 natten/flake.lock create mode 100644 natten/flake.nix create mode 100644 natten/scripts/autogen_blackwell_fmha.py create mode 100644 natten/scripts/autogen_blackwell_fmha_bwd.py create mode 100644 natten/scripts/autogen_blackwell_fna.py create mode 100644 natten/scripts/autogen_blackwell_fna_bwd.py create mode 100644 natten/scripts/autogen_fmha.py create mode 100644 natten/scripts/autogen_fna.py create mode 100644 natten/scripts/autogen_hopper_fmha.py create mode 100644 natten/scripts/autogen_hopper_fmha_bwd.py create mode 100644 natten/scripts/autogen_hopper_fna.py create mode 100644 natten/scripts/autogen_hopper_fna_bwd.py create mode 100644 natten/scripts/autogen_reference_fna.py create mode 100644 natten/scripts/generate_build_toml.py create mode 100755 natten/scripts/regen.sh create mode 100644 natten/tests/__init__.py create mode 100644 natten/tests/conftest.py create mode 100644 natten/tests/test_attn_merge.py create mode 100644 natten/tests/test_blackwell_fna.py create mode 100644 natten/tests/test_compute_delta.py create mode 100644 natten/tests/test_fmha.py create mode 100644 natten/tests/test_fmha_varlen.py create mode 100644 natten/tests/test_fna.py create mode 100644 natten/tests/test_hopper_fna.py create mode 100644 natten/tests/test_kernels_ci.py create mode 100644 natten/tests/test_token_permute.py create mode 100644 natten/tests/test_torch_compile.py create mode 100644 natten/tests/utils.py create mode 100644 natten/torch-ext/natten/__init__.py create mode 100644 natten/torch-ext/natten/_environment.py create mode 100644 natten/torch-ext/natten/_libnatten/__init__.py create mode 100644 natten/torch-ext/natten/_libnatten/torch_wrappers.py create mode 100644 natten/torch-ext/natten/attn_merge.py create mode 100644 natten/torch-ext/natten/backends/__init__.py create mode 100644 natten/torch-ext/natten/backends/blackwell_fmha.py create mode 100644 natten/torch-ext/natten/backends/blackwell_fna.py create mode 100644 natten/torch-ext/natten/backends/configs/__init__.py create mode 100644 natten/torch-ext/natten/backends/configs/checks.py create mode 100644 natten/torch-ext/natten/backends/configs/cutlass/__init__.py create mode 100644 natten/torch-ext/natten/backends/configs/cutlass/backward_knobs.py create mode 100644 natten/torch-ext/natten/backends/configs/cutlass/fna_backward_128x128.py create mode 100644 natten/torch-ext/natten/backends/configs/cutlass/fna_backward_128x64.py create mode 100644 natten/torch-ext/natten/backends/configs/cutlass/fna_backward_64x64.py create mode 100644 natten/torch-ext/natten/backends/configs/cutlass/fna_forward_32x128.py create mode 100644 natten/torch-ext/natten/backends/configs/cutlass/fna_forward_64x128.py create mode 100644 natten/torch-ext/natten/backends/configs/cutlass/fna_forward_64x64.py create mode 100644 natten/torch-ext/natten/backends/configs/cutlass_blackwell/__init__.py create mode 100644 natten/torch-ext/natten/backends/configs/cutlass_hopper/__init__.py create mode 100644 natten/torch-ext/natten/backends/configs/flex/__init__.py create mode 100644 natten/torch-ext/natten/backends/flex.py create mode 100644 natten/torch-ext/natten/backends/fmha.py create mode 100644 natten/torch-ext/natten/backends/fna.py create mode 100644 natten/torch-ext/natten/backends/hopper_fmha.py create mode 100644 natten/torch-ext/natten/backends/hopper_fna.py create mode 100644 natten/torch-ext/natten/backends/reference.py create mode 100644 natten/torch-ext/natten/context.py create mode 100644 natten/torch-ext/natten/functional.py create mode 100644 natten/torch-ext/natten/modules.py create mode 100644 natten/torch-ext/natten/token_permute/__init__.py create mode 100644 natten/torch-ext/natten/token_permute/cutlass_impl.py create mode 100644 natten/torch-ext/natten/token_permute/frontend.py create mode 100644 natten/torch-ext/natten/token_permute/torch_impl.py create mode 100644 natten/torch-ext/natten/types.py create mode 100644 natten/torch-ext/natten/utils/__init__.py create mode 100644 natten/torch-ext/natten/utils/checks.py create mode 100644 natten/torch-ext/natten/utils/device.py create mode 100644 natten/torch-ext/natten/utils/dtype.py create mode 100644 natten/torch-ext/natten/utils/environment.py create mode 100644 natten/torch-ext/natten/utils/log.py create mode 100644 natten/torch-ext/natten/utils/tensor.py create mode 100644 natten/torch-ext/natten/utils/testing.py create mode 100644 natten/torch-ext/natten/utils/tuples.py create mode 100644 natten/torch-ext/natten/utils/varlen.py create mode 100644 natten/torch-ext/natten/version.py create mode 100644 natten/torch-ext/torch_binding.cpp diff --git a/natten/CARD.md b/natten/CARD.md new file mode 100644 index 00000000..60f8ab8d --- /dev/null +++ b/natten/CARD.md @@ -0,0 +1,112 @@ +--- +license: mit +tags: +- kernel +--- + +![Status](https://hubwebhook.dholtz.com/shield?repo=kernels-community/natten) + +# NATTEN — Neighborhood Attention kernels + +CUDA kernels for [NATTEN](https://natten.org) (Neighborhood Attention +Extension, [SHI-Labs/NATTEN](https://github.com/SHI-Labs/NATTEN)), packaged +for the [kernels](https://github.com/huggingface/kernels) library. Vendored +from NATTEN v0.21.6 (MIT license). + +Neighborhood attention restricts each token's attention to a sliding local +window over 1-D, 2-D, or 3-D token layouts, with optional stride, dilation, +and per-dimension causal masking. This kernel ships NATTEN's fused +neighborhood attention (FNA) and FMHA kernels along with the thin functional +frontend (`na1d`, `na2d`, `na3d`, `attention`, `merge_attentions`) and its +backend/config selection logic. + +## Usage + +```python +import torch +from kernels import get_kernel + +natten = get_kernel("kernels-community/natten", version=1) + +# 2-D neighborhood attention over a 32x32 token layout: +# [batch, *token_layout, heads, head_dim] +q = torch.randn(1, 32, 32, 8, 64, device="cuda", dtype=torch.bfloat16) +k = torch.randn_like(q) +v = torch.randn_like(q) + +out = natten.functional.na2d(q, k, v, kernel_size=(7, 7)) + +# 3-D, with stride/dilation/causal masking per dim: +q3 = torch.randn(1, 8, 16, 16, 8, 64, device="cuda", dtype=torch.bfloat16) +k3, v3 = torch.randn_like(q3), torch.randn_like(q3) +out3 = natten.functional.na3d( + q3, k3, v3, kernel_size=(3, 5, 5), dilation=(1, 2, 2), is_causal=(True, False, False) +) +``` + +All functions are differentiable and torch.compile-compatible (the underlying +ops are registered via `TORCH_LIBRARY` with fake/meta implementations). + +## Backend / architecture support + +| Backend | Compute capabilities | Notes | +|---|---|---| +| Reference CUDA kernels | 7.0 – 12.0 | correctness fallback, all dims | +| CUTLASS 2.X FNA / FMHA (`cutlass-fna`) | 7.0 – 12.0 (SM80 tensor-core paths) | default on pre-Hopper | +| Hopper FNA / FMHA (`hopper-fna`) | 9.0a | H100/H200 | +| Blackwell FNA / FMHA (`blackwell-fna`) | 10.0a | B200/GB200; requires CUDA ≥ 12.8, so cu126 builds fall back to `cutlass-fna` / reference on these GPUs | +| Flex Attention backend (`flex-fna`) | any | pure PyTorch, via `torch.nn.attention.flex_attention` | + +Known gaps vs. upstream wheels: + +- No SM103 (B300 / Blackwell Ultra) binaries — kernel-builder does not + currently target `10.3a`. +- SM120 (RTX 50 series) gets the CUTLASS 2.X and reference paths only, same + as upstream (NATTEN has no SM120-specific fused kernels). +- No ROCm or CPU support (upstream dropped these in v0.20). + +The backend is chosen automatically per GPU; pass `backend="..."` to +`na1d`/`na2d`/`na3d` to override. + +## Developing + +The ~144 kernel instantiation TUs under `csrc/autogen/` are committed +codegen output. To regenerate (e.g. when bumping the vendored NATTEN +version), run: + +```bash +scripts/regen.sh # re-runs autogen + rewrites build.toml src lists +``` + +Build and test with kernel-builder: + +```bash +nix run .#build-and-copy -L +nix run .#ci-test -L +``` + +## Credits + +All kernels and the Python frontend are by Ali Hassani and the NATTEN +contributors ([SHI-Labs/NATTEN](https://github.com/SHI-Labs/NATTEN), MIT). +This repository only repackages them for the Kernel Hub: pybind11 bindings +are replaced with `TORCH_LIBRARY` registrations (Python limited API), and +build-time codegen is committed statically. + +If you use NATTEN, please cite: + +```bibtex +@inproceedings{hassani2023neighborhood, + title = {Neighborhood Attention Transformer}, + author = {Ali Hassani and Steven Walton and Jiachen Li and Shen Li and Humphrey Shi}, + year = 2023, + booktitle = {IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)}, +} +@misc{hassani2024faster, + title = {Faster Neighborhood Attention: Reducing the O(n^2) Cost of Self Attention at the Threadblock Level}, + author = {Ali Hassani and Wen-Mei Hwu and Humphrey Shi}, + year = 2024, + eprint = {2403.04690}, + archivePrefix = {arXiv}, +} +``` diff --git a/natten/LICENSE b/natten/LICENSE new file mode 100644 index 00000000..d47c679f --- /dev/null +++ b/natten/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2022 - 2026 Ali Hassani. + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/natten/benchmarks/benchmark.py b/natten/benchmarks/benchmark.py new file mode 100644 index 00000000..58fcc653 --- /dev/null +++ b/natten/benchmarks/benchmark.py @@ -0,0 +1,73 @@ +# Micro-benchmark for the fused neighborhood attention path. +# +# Compares NATTEN's auto-selected fused backend against PyTorch SDPA over the +# full sequence (which computes strictly more attention, but is the baseline +# people care about) across a few window sizes. +# +# Run inside `kernel-builder devshell` / `testshell`, or any env where the +# built `natten` package is importable. + +import argparse +import time + +import torch +from torch.nn.functional import scaled_dot_product_attention + +from natten.functional import na2d + + +def benchmark(fn, warmup: int = 10, iters: int = 50) -> float: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start = time.perf_counter() + for _ in range(iters): + fn() + torch.cuda.synchronize() + return (time.perf_counter() - start) / iters * 1e3 + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--batch", type=int, default=1) + parser.add_argument("--heads", type=int, default=8) + parser.add_argument("--head-dim", type=int, default=64) + parser.add_argument("--size", type=int, default=64, help="2-D token layout side") + parser.add_argument( + "--dtype", choices=["float16", "bfloat16"], default="bfloat16" + ) + args = parser.parse_args() + + dtype = getattr(torch, args.dtype) + device = "cuda" + shape = (args.batch, args.size, args.size, args.heads, args.head_dim) + + q = torch.randn(shape, device=device, dtype=dtype) + k = torch.randn_like(q) + v = torch.randn_like(q) + + # [B, H, seq, D] view for SDPA + q_sdpa = q.flatten(1, 2).permute(0, 2, 1, 3).contiguous() + k_sdpa = k.flatten(1, 2).permute(0, 2, 1, 3).contiguous() + v_sdpa = v.flatten(1, 2).permute(0, 2, 1, 3).contiguous() + + sdpa_ms = benchmark(lambda: scaled_dot_product_attention(q_sdpa, k_sdpa, v_sdpa)) + print( + f"device={torch.cuda.get_device_name()} dtype={args.dtype} " + f"layout={args.size}x{args.size} heads={args.heads} head_dim={args.head_dim}" + ) + print(f"{'kernel':>24} {'ms':>10} {'vs SDPA':>10}") + print(f"{'sdpa (full self-attn)':>24} {sdpa_ms:>10.3f} {'1.00x':>10}") + + for window in (7, 13, 21, 33): + if window > args.size: + continue + na_ms = benchmark(lambda: na2d(q, k, v, kernel_size=(window, window))) + print( + f"{f'na2d k={window}x{window}':>24} {na_ms:>10.3f} " + f"{f'{sdpa_ms / na_ms:.2f}x':>10}" + ) + + +if __name__ == "__main__": + main() diff --git a/natten/build.toml b/natten/build.toml new file mode 100644 index 00000000..22934722 --- /dev/null +++ b/natten/build.toml @@ -0,0 +1,459 @@ +[general] +name = "natten" +version = 1 +edition = 5 +license = "MIT" +upstream = "https://github.com/SHI-Labs/NATTEN" +backends = ["cuda"] + +[general.hub] +repo-id = "kernels-community/natten" + +[torch] +src = ["torch-ext/torch_binding.cpp"] +include = ["csrc/include"] + +[kernel.natten] +backend = "cuda" +cuda-capabilities = [ + "7.0", + "7.5", + "8.0", + "8.6", + "8.9", + "9.0", + "10.0", + "12.0", +] +cuda-flags = [ + "-O3", + "-std=c++17", + "--expt-relaxed-constexpr", + "--extended-lambda", + "--use_fast_math", + "--ftemplate-backtrace-limit=0", + "-Xfatbin=-compress-all", + "-DNDEBUG", + "-DNATTEN_WITH_CUTLASS", + "-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1", +] +depends = [ + "torch", + "cutlass_4_5", +] +include = [ + "csrc/include", + "csrc/autogen/include", +] +src = [ + "csrc/include/natten/blackwell_fmha.h", + "csrc/include/natten/blackwell_fna.h", + "csrc/include/natten/compute_delta.h", + "csrc/include/natten/cuda/fmha/epilogue/epilogue_pipelined.h", + "csrc/include/natten/cuda/fmha/epilogue/epilogue_rescale_output.h", + "csrc/include/natten/cuda/fmha/epilogue/epilogue_thread_apply_logsumexp.h", + "csrc/include/natten/cuda/fmha/fmha_backward.cuh", + "csrc/include/natten/cuda/fmha/fmha_forward.cuh", + "csrc/include/natten/cuda/fmha/gemm/custom_mma.h", + "csrc/include/natten/cuda/fmha/gemm/custom_mma_base.h", + "csrc/include/natten/cuda/fmha/gemm/custom_mma_multistage.h", + "csrc/include/natten/cuda/fmha/gemm/custom_mma_pipelined.h", + "csrc/include/natten/cuda/fmha/gemm/find_default_mma.h", + "csrc/include/natten/cuda/fmha/gemm/mma_accum_lambda_iterator.h", + "csrc/include/natten/cuda/fmha/gemm/mma_from_smem.h", + "csrc/include/natten/cuda/fmha/gemm_kernel_utils.h", + "csrc/include/natten/cuda/fmha/iterators/default_warp_iterator_from_smem.h", + "csrc/include/natten/cuda/fmha/iterators/epilogue_predicated_tile_iterator.h", + "csrc/include/natten/cuda/fmha/iterators/make_residual_last.h", + "csrc/include/natten/cuda/fmha/iterators/predicated_tile_access_iterator_residual_last.h", + "csrc/include/natten/cuda/fmha/iterators/predicated_tile_iterator_residual_last.h", + "csrc/include/natten/cuda/fmha/iterators/transpose_warp_iterator.h", + "csrc/include/natten/cuda/fmha/iterators/warp_iterator_from_smem.h", + "csrc/include/natten/cuda/fmha/kernel_backward.h", + "csrc/include/natten/cuda/fmha/kernel_forward.h", + "csrc/include/natten/cuda/fmha/transform/tile_smem_loader.h", + "csrc/include/natten/cuda/fmha_blackwell/collective/fmha_common.hpp", + "csrc/include/natten/cuda/fmha_blackwell/collective/fmha_fusion.hpp", + "csrc/include/natten/cuda/fmha_blackwell/collective/sm100_fmha_fwd_epilogue_tma_warpspecialized.hpp", + "csrc/include/natten/cuda/fmha_blackwell/collective/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp", + "csrc/include/natten/cuda/fmha_blackwell/collective/sm100_fmha_load_tma_warpspecialized.hpp", + "csrc/include/natten/cuda/fmha_blackwell/common/pow_2.hpp", + "csrc/include/natten/cuda/fmha_blackwell/device/fmha_bwd_sm100.hpp", + "csrc/include/natten/cuda/fmha_blackwell/device/fmha_sm100.hpp", + "csrc/include/natten/cuda/fmha_blackwell/fmha_backward.cuh", + "csrc/include/natten/cuda/fmha_blackwell/fmha_forward.cuh", + "csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_causal_tile_scheduler.hpp", + "csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_kernel_bwd_convert.hpp", + "csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_kernel_bwd_sum_OdO.hpp", + "csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_tile_scheduler.hpp", + "csrc/include/natten/cuda/fmha_blackwell/kernel/sm100_fmha_bwd_kernel_tma_warpspecialized.hpp", + "csrc/include/natten/cuda/fmha_blackwell/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp", + "csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_bwd_tma_warpspecialized.hpp", + "csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_load.hpp", + "csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_softmax.hpp", + "csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_tma.hpp", + "csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_tma_warpspecialized.hpp", + "csrc/include/natten/cuda/fmha_hopper/collective/fmha_common.hpp", + "csrc/include/natten/cuda/fmha_hopper/collective/fmha_epilogue.hpp", + "csrc/include/natten/cuda/fmha_hopper/collective/fmha_epilogue_bwd.hpp", + "csrc/include/natten/cuda/fmha_hopper/collective/fmha_fusion.hpp", + "csrc/include/natten/cuda/fmha_hopper/collective/fmha_varlen.hpp", + "csrc/include/natten/cuda/fmha_hopper/device/fmha_bwd_sm90.hpp", + "csrc/include/natten/cuda/fmha_hopper/device/fmha_sm90.hpp", + "csrc/include/natten/cuda/fmha_hopper/fmha_backward.cuh", + "csrc/include/natten/cuda/fmha_hopper/fmha_forward.cuh", + "csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_builder.hpp", + "csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_bwd_convert.hpp", + "csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_bwd_sum_OdO.hpp", + "csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_tma.hpp", + "csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_tma_warpspecialized.hpp", + "csrc/include/natten/cuda/fmha_hopper/kernel/fmha_options.hpp", + "csrc/include/natten/cuda/fmha_hopper/kernel/fmha_tile_scheduler.hpp", + "csrc/include/natten/cuda/fna/epilogue/epilogue_pipelined.h", + "csrc/include/natten/cuda/fna/epilogue/epilogue_rescale_output.h", + "csrc/include/natten/cuda/fna/epilogue/epilogue_thread_apply_logsumexp.h", + "csrc/include/natten/cuda/fna/epilogue/predicated_tile_iterator.h", + "csrc/include/natten/cuda/fna/epilogue/predicated_tile_iterator_params.h", + "csrc/include/natten/cuda/fna/fna_backward.cuh", + "csrc/include/natten/cuda/fna/fna_forward.cuh", + "csrc/include/natten/cuda/fna/gemm/custom_mma.h", + "csrc/include/natten/cuda/fna/gemm/custom_mma_base.h", + "csrc/include/natten/cuda/fna/gemm/custom_mma_multistage.h", + "csrc/include/natten/cuda/fna/gemm/custom_mma_pipelined.h", + "csrc/include/natten/cuda/fna/gemm/find_default_mma.h", + "csrc/include/natten/cuda/fna/gemm/mma_accum_lambda_iterator.h", + "csrc/include/natten/cuda/fna/gemm/mma_from_smem.h", + "csrc/include/natten/cuda/fna/gemm/replace_mma_iterators.h", + "csrc/include/natten/cuda/fna/gemm_kernel_utils.h", + "csrc/include/natten/cuda/fna/iterators/default_warp_iterator_from_smem.h", + "csrc/include/natten/cuda/fna/iterators/epilogue_predicated_tile_iterator.h", + "csrc/include/natten/cuda/fna/iterators/make_residual_last.h", + "csrc/include/natten/cuda/fna/iterators/predicated_tile_access_iterator.h", + "csrc/include/natten/cuda/fna/iterators/predicated_tile_access_iterator_residual_last.h", + "csrc/include/natten/cuda/fna/iterators/predicated_tile_iterator.h", + "csrc/include/natten/cuda/fna/iterators/predicated_tile_iterator_residual_last.h", + "csrc/include/natten/cuda/fna/iterators/transpose_warp_iterator.h", + "csrc/include/natten/cuda/fna/iterators/warp_iterator_from_smem.h", + "csrc/include/natten/cuda/fna/kernel_backward.h", + "csrc/include/natten/cuda/fna/kernel_forward.h", + "csrc/include/natten/cuda/fna/na_utils.cuh", + "csrc/include/natten/cuda/fna_blackwell/collective/fna_common.hpp", + "csrc/include/natten/cuda/fna_blackwell/collective/fna_fusion.hpp", + "csrc/include/natten/cuda/fna_blackwell/collective/fna_fusion_bwd.hpp", + "csrc/include/natten/cuda/fna_blackwell/collective/sm100_fna_fwd_epilogue_tma_warpspecialized.hpp", + "csrc/include/natten/cuda/fna_blackwell/collective/sm100_fna_fwd_mainloop_tma_warpspecialized.hpp", + "csrc/include/natten/cuda/fna_blackwell/collective/sm100_fna_load_tma_warpspecialized.hpp", + "csrc/include/natten/cuda/fna_blackwell/common/pow_2.hpp", + "csrc/include/natten/cuda/fna_blackwell/device/fna_bwd_sm100.hpp", + "csrc/include/natten/cuda/fna_blackwell/device/fna_sm100.hpp", + "csrc/include/natten/cuda/fna_blackwell/fna_backward.cuh", + "csrc/include/natten/cuda/fna_blackwell/fna_forward.cuh", + "csrc/include/natten/cuda/fna_blackwell/kernel/sm100_fna_bwd_kernel_tma_warpspecialized.hpp", + "csrc/include/natten/cuda/fna_blackwell/kernel/sm100_fna_fwd_kernel_tma_warpspecialized.hpp", + "csrc/include/natten/cuda/fna_hopper/collective/fna_collective_bwd_tma_warpspecialized.hpp", + "csrc/include/natten/cuda/fna_hopper/collective/fna_collective_load.hpp", + "csrc/include/natten/cuda/fna_hopper/collective/fna_collective_softmax.hpp", + "csrc/include/natten/cuda/fna_hopper/collective/fna_collective_tma.hpp", + "csrc/include/natten/cuda/fna_hopper/collective/fna_collective_tma_warpspecialized.hpp", + "csrc/include/natten/cuda/fna_hopper/collective/fna_common.hpp", + "csrc/include/natten/cuda/fna_hopper/collective/fna_fusion.hpp", + "csrc/include/natten/cuda/fna_hopper/collective/fna_fusion_bwd.hpp", + "csrc/include/natten/cuda/fna_hopper/device/fna_bwd_sm90.hpp", + "csrc/include/natten/cuda/fna_hopper/device/fna_sm90.hpp", + "csrc/include/natten/cuda/fna_hopper/fna_backward.cuh", + "csrc/include/natten/cuda/fna_hopper/fna_forward.cuh", + "csrc/include/natten/cuda/fna_hopper/kernel/fna_kernel_builder.hpp", + "csrc/include/natten/cuda/hopper_fmha_fna.h", + "csrc/include/natten/cuda/reduction/compute_delta.cuh", + "csrc/include/natten/cuda/reduction/fmha_kernel_bwd_sum_OdO.hpp", + "csrc/include/natten/cuda/reference/fna_reference_backward.hpp", + "csrc/include/natten/cuda/reference/fna_reference_forward.hpp", + "csrc/include/natten/cuda/reference/mask.hpp", + "csrc/include/natten/cuda/reference/utils.hpp", + "csrc/include/natten/cuda/tokperm/layouts.hpp", + "csrc/include/natten/cuda/tokperm/token_permute_kernel.cuh", + "csrc/include/natten/cuda/tokperm/tokperm.hpp", + "csrc/include/natten/cuda/tokperm/utils/permute.cuh", + "csrc/include/natten/cuda/tokperm/utils/stride.cuh", + "csrc/include/natten/cuda/tokperm/utils/tuple.cuh", + "csrc/include/natten/cuda/utils/cuda.h", + "csrc/include/natten/cuda/utils/cutlass.cuh", + "csrc/include/natten/cuda/utils/generic_cutlass_device.hpp", + "csrc/include/natten/fmha.h", + "csrc/include/natten/fna.h", + "csrc/include/natten/helpers.h", + "csrc/include/natten/hopper_fmha.h", + "csrc/include/natten/hopper_fna.h", + "csrc/include/natten/natten.h", + "csrc/include/natten/reference.h", + "csrc/include/natten/token_permute.h", + "csrc/src/blackwell_fmha.cu", + "csrc/src/blackwell_fna_backward.cu", + "csrc/src/blackwell_fna_forward.cu", + "csrc/src/compute_delta.cu", + "csrc/src/fmha.cu", + "csrc/src/fna_backward.cu", + "csrc/src/fna_forward.cu", + "csrc/src/hopper_fmha.cu", + "csrc/src/hopper_fna_backward.cu", + "csrc/src/hopper_fna_forward.cu", + "csrc/src/reference_backward.cu", + "csrc/src/reference_forward.cu", + "csrc/src/token_permute.cu", + "csrc/src/token_unpermute.cu", + "csrc/autogen/include/natten_autogen/cuda/reference/dispatch_cm.h", + "csrc/autogen/include/natten_autogen/cuda/reference/dispatch_dtype.h", + "csrc/autogen/include/natten_autogen/cuda/reference/interface.h", + "csrc/autogen/include/natten_autogen/cuda/reference/kernels.h", + "csrc/autogen/src/cuda/reference/source_0.cu", + "csrc/autogen/src/cuda/reference/source_1.cu", + "csrc/autogen/include/natten_autogen/cuda/fna/dispatch_cm.h", + "csrc/autogen/include/natten_autogen/cuda/fna/dispatch_device.h", + "csrc/autogen/include/natten_autogen/cuda/fna/dispatch_dtype.h", + "csrc/autogen/include/natten_autogen/cuda/fna/interface.h", + "csrc/autogen/include/natten_autogen/cuda/fna/kernels.h", + "csrc/autogen/src/cuda/fna/source_0.cu", + "csrc/autogen/src/cuda/fna/source_1.cu", + "csrc/autogen/src/cuda/fna/source_10.cu", + "csrc/autogen/src/cuda/fna/source_11.cu", + "csrc/autogen/src/cuda/fna/source_12.cu", + "csrc/autogen/src/cuda/fna/source_13.cu", + "csrc/autogen/src/cuda/fna/source_14.cu", + "csrc/autogen/src/cuda/fna/source_15.cu", + "csrc/autogen/src/cuda/fna/source_16.cu", + "csrc/autogen/src/cuda/fna/source_17.cu", + "csrc/autogen/src/cuda/fna/source_18.cu", + "csrc/autogen/src/cuda/fna/source_19.cu", + "csrc/autogen/src/cuda/fna/source_2.cu", + "csrc/autogen/src/cuda/fna/source_20.cu", + "csrc/autogen/src/cuda/fna/source_21.cu", + "csrc/autogen/src/cuda/fna/source_22.cu", + "csrc/autogen/src/cuda/fna/source_23.cu", + "csrc/autogen/src/cuda/fna/source_24.cu", + "csrc/autogen/src/cuda/fna/source_25.cu", + "csrc/autogen/src/cuda/fna/source_26.cu", + "csrc/autogen/src/cuda/fna/source_27.cu", + "csrc/autogen/src/cuda/fna/source_28.cu", + "csrc/autogen/src/cuda/fna/source_29.cu", + "csrc/autogen/src/cuda/fna/source_3.cu", + "csrc/autogen/src/cuda/fna/source_30.cu", + "csrc/autogen/src/cuda/fna/source_31.cu", + "csrc/autogen/src/cuda/fna/source_32.cu", + "csrc/autogen/src/cuda/fna/source_33.cu", + "csrc/autogen/src/cuda/fna/source_34.cu", + "csrc/autogen/src/cuda/fna/source_35.cu", + "csrc/autogen/src/cuda/fna/source_36.cu", + "csrc/autogen/src/cuda/fna/source_37.cu", + "csrc/autogen/src/cuda/fna/source_38.cu", + "csrc/autogen/src/cuda/fna/source_39.cu", + "csrc/autogen/src/cuda/fna/source_4.cu", + "csrc/autogen/src/cuda/fna/source_40.cu", + "csrc/autogen/src/cuda/fna/source_41.cu", + "csrc/autogen/src/cuda/fna/source_42.cu", + "csrc/autogen/src/cuda/fna/source_43.cu", + "csrc/autogen/src/cuda/fna/source_44.cu", + "csrc/autogen/src/cuda/fna/source_45.cu", + "csrc/autogen/src/cuda/fna/source_46.cu", + "csrc/autogen/src/cuda/fna/source_47.cu", + "csrc/autogen/src/cuda/fna/source_48.cu", + "csrc/autogen/src/cuda/fna/source_49.cu", + "csrc/autogen/src/cuda/fna/source_5.cu", + "csrc/autogen/src/cuda/fna/source_50.cu", + "csrc/autogen/src/cuda/fna/source_51.cu", + "csrc/autogen/src/cuda/fna/source_52.cu", + "csrc/autogen/src/cuda/fna/source_53.cu", + "csrc/autogen/src/cuda/fna/source_54.cu", + "csrc/autogen/src/cuda/fna/source_55.cu", + "csrc/autogen/src/cuda/fna/source_56.cu", + "csrc/autogen/src/cuda/fna/source_57.cu", + "csrc/autogen/src/cuda/fna/source_58.cu", + "csrc/autogen/src/cuda/fna/source_59.cu", + "csrc/autogen/src/cuda/fna/source_6.cu", + "csrc/autogen/src/cuda/fna/source_60.cu", + "csrc/autogen/src/cuda/fna/source_61.cu", + "csrc/autogen/src/cuda/fna/source_62.cu", + "csrc/autogen/src/cuda/fna/source_63.cu", + "csrc/autogen/src/cuda/fna/source_7.cu", + "csrc/autogen/src/cuda/fna/source_8.cu", + "csrc/autogen/src/cuda/fna/source_9.cu", + "csrc/autogen/include/natten_autogen/cuda/fmha/dispatch_dtype.h", + "csrc/autogen/include/natten_autogen/cuda/fmha/interface.h", + "csrc/autogen/include/natten_autogen/cuda/fmha/kernels.h", + "csrc/autogen/src/cuda/fmha/source_0.cu", + "csrc/autogen/src/cuda/fmha/source_1.cu", + "csrc/autogen/src/cuda/fmha/source_2.cu", + "csrc/autogen/src/cuda/fmha/source_3.cu", + "csrc/autogen/src/cuda/fmha/source_4.cu", + "csrc/autogen/src/cuda/fmha/source_5.cu", +] + +[kernel.natten_hopper] +backend = "cuda" +cuda-capabilities = [ + "9.0a", +] +cuda-flags = [ + "-O3", + "-std=c++17", + "--expt-relaxed-constexpr", + "--extended-lambda", + "--use_fast_math", + "--ftemplate-backtrace-limit=0", + "-Xfatbin=-compress-all", + "-DNDEBUG", + "-DNATTEN_WITH_CUTLASS", + "-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1", + "-DNATTEN_WITH_HOPPER_FNA=1", +] +depends = [ + "torch", + "cutlass_4_5", +] +include = [ + "csrc/include", + "csrc/autogen/include", +] +src = [ + "csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_cm.h", + "csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_dtype.h", + "csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_head_dim.h", + "csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_tile_shape.h", + "csrc/autogen/include/natten_autogen/cuda/hopper_fna/interface.h", + "csrc/autogen/include/natten_autogen/cuda/hopper_fna/kernels.h", + "csrc/autogen/src/cuda/hopper_fna/source_0.cu", + "csrc/autogen/src/cuda/hopper_fna/source_1.cu", + "csrc/autogen/src/cuda/hopper_fna/source_2.cu", + "csrc/autogen/src/cuda/hopper_fna/source_3.cu", + "csrc/autogen/src/cuda/hopper_fna/source_4.cu", + "csrc/autogen/src/cuda/hopper_fna/source_5.cu", + "csrc/autogen/src/cuda/hopper_fna/source_6.cu", + "csrc/autogen/src/cuda/hopper_fna/source_7.cu", + "csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_cm.h", + "csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_dtype.h", + "csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_head_dim.h", + "csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_tile_shape.h", + "csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/interface.h", + "csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/kernels.h", + "csrc/autogen/src/cuda/hopper_fna_bwd/source_0.cu", + "csrc/autogen/src/cuda/hopper_fna_bwd/source_1.cu", + "csrc/autogen/src/cuda/hopper_fna_bwd/source_2.cu", + "csrc/autogen/src/cuda/hopper_fna_bwd/source_3.cu", + "csrc/autogen/include/natten_autogen/cuda/hopper_fmha/dispatch_head_dim.h", + "csrc/autogen/include/natten_autogen/cuda/hopper_fmha/dispatch_tile_size.h", + "csrc/autogen/include/natten_autogen/cuda/hopper_fmha/interface.h", + "csrc/autogen/include/natten_autogen/cuda/hopper_fmha/kernels.h", + "csrc/autogen/src/cuda/hopper_fmha/source_0.cu", + "csrc/autogen/src/cuda/hopper_fmha/source_1.cu", + "csrc/autogen/src/cuda/hopper_fmha/source_2.cu", + "csrc/autogen/src/cuda/hopper_fmha/source_3.cu", + "csrc/autogen/src/cuda/hopper_fmha/source_4.cu", + "csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/dispatch_head_dim.h", + "csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/dispatch_tile_size.h", + "csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/interface.h", + "csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/kernels.h", + "csrc/autogen/src/cuda/hopper_fmha_bwd/source_0.cu", + "csrc/autogen/src/cuda/hopper_fmha_bwd/source_1.cu", + "csrc/autogen/src/cuda/hopper_fmha_bwd/source_2.cu", + "csrc/autogen/src/cuda/hopper_fmha_bwd/source_3.cu", + "csrc/autogen/src/cuda/hopper_fmha_bwd/source_4.cu", +] + +[kernel.natten_blackwell] +backend = "cuda" +cuda-minver = "12.8" +cuda-capabilities = [ + "10.0a", +] +cuda-flags = [ + "-O3", + "-std=c++17", + "--expt-relaxed-constexpr", + "--extended-lambda", + "--use_fast_math", + "--ftemplate-backtrace-limit=0", + "-Xfatbin=-compress-all", + "-DNDEBUG", + "-DNATTEN_WITH_CUTLASS", + "-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1", + "-DNATTEN_WITH_BLACKWELL_FNA=1", +] +depends = [ + "torch", + "cutlass_4_5", +] +include = [ + "csrc/include", + "csrc/autogen/include", +] +src = [ + "csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_cm.h", + "csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_dtype.h", + "csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_head_dim.h", + "csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_tile_shape.h", + "csrc/autogen/include/natten_autogen/cuda/blackwell_fna/interface.h", + "csrc/autogen/include/natten_autogen/cuda/blackwell_fna/kernels.h", + "csrc/autogen/src/cuda/blackwell_fna/source_0.cu", + "csrc/autogen/src/cuda/blackwell_fna/source_1.cu", + "csrc/autogen/src/cuda/blackwell_fna/source_10.cu", + "csrc/autogen/src/cuda/blackwell_fna/source_11.cu", + "csrc/autogen/src/cuda/blackwell_fna/source_12.cu", + "csrc/autogen/src/cuda/blackwell_fna/source_13.cu", + "csrc/autogen/src/cuda/blackwell_fna/source_14.cu", + "csrc/autogen/src/cuda/blackwell_fna/source_15.cu", + "csrc/autogen/src/cuda/blackwell_fna/source_16.cu", + "csrc/autogen/src/cuda/blackwell_fna/source_17.cu", + "csrc/autogen/src/cuda/blackwell_fna/source_18.cu", + "csrc/autogen/src/cuda/blackwell_fna/source_19.cu", + "csrc/autogen/src/cuda/blackwell_fna/source_2.cu", + "csrc/autogen/src/cuda/blackwell_fna/source_20.cu", + "csrc/autogen/src/cuda/blackwell_fna/source_21.cu", + "csrc/autogen/src/cuda/blackwell_fna/source_22.cu", + "csrc/autogen/src/cuda/blackwell_fna/source_23.cu", + "csrc/autogen/src/cuda/blackwell_fna/source_24.cu", + "csrc/autogen/src/cuda/blackwell_fna/source_25.cu", + "csrc/autogen/src/cuda/blackwell_fna/source_26.cu", + "csrc/autogen/src/cuda/blackwell_fna/source_27.cu", + "csrc/autogen/src/cuda/blackwell_fna/source_3.cu", + "csrc/autogen/src/cuda/blackwell_fna/source_4.cu", + "csrc/autogen/src/cuda/blackwell_fna/source_5.cu", + "csrc/autogen/src/cuda/blackwell_fna/source_6.cu", + "csrc/autogen/src/cuda/blackwell_fna/source_7.cu", + "csrc/autogen/src/cuda/blackwell_fna/source_8.cu", + "csrc/autogen/src/cuda/blackwell_fna/source_9.cu", + "csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_cm.h", + "csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_dtype.h", + "csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_head_dim.h", + "csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_tile_shape.h", + "csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/interface.h", + "csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/kernels.h", + "csrc/autogen/src/cuda/blackwell_fna_bwd/source_0.cu", + "csrc/autogen/src/cuda/blackwell_fna_bwd/source_1.cu", + "csrc/autogen/src/cuda/blackwell_fna_bwd/source_10.cu", + "csrc/autogen/src/cuda/blackwell_fna_bwd/source_11.cu", + "csrc/autogen/src/cuda/blackwell_fna_bwd/source_12.cu", + "csrc/autogen/src/cuda/blackwell_fna_bwd/source_13.cu", + "csrc/autogen/src/cuda/blackwell_fna_bwd/source_2.cu", + "csrc/autogen/src/cuda/blackwell_fna_bwd/source_3.cu", + "csrc/autogen/src/cuda/blackwell_fna_bwd/source_4.cu", + "csrc/autogen/src/cuda/blackwell_fna_bwd/source_5.cu", + "csrc/autogen/src/cuda/blackwell_fna_bwd/source_6.cu", + "csrc/autogen/src/cuda/blackwell_fna_bwd/source_7.cu", + "csrc/autogen/src/cuda/blackwell_fna_bwd/source_8.cu", + "csrc/autogen/src/cuda/blackwell_fna_bwd/source_9.cu", + "csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/dispatch_head_dim.h", + "csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/dispatch_tile_size.h", + "csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/interface.h", + "csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/kernels.h", + "csrc/autogen/src/cuda/blackwell_fmha/source_0.cu", + "csrc/autogen/src/cuda/blackwell_fmha/source_1.cu", + "csrc/autogen/src/cuda/blackwell_fmha/source_2.cu", + "csrc/autogen/src/cuda/blackwell_fmha/source_3.cu", + "csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/dispatch_head_dim.h", + "csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/dispatch_tile_size.h", + "csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/interface.h", + "csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/kernels.h", + "csrc/autogen/src/cuda/blackwell_fmha_bwd/source_0.cu", + "csrc/autogen/src/cuda/blackwell_fmha_bwd/source_1.cu", + "csrc/autogen/src/cuda/blackwell_fmha_bwd/source_2.cu", + "csrc/autogen/src/cuda/blackwell_fmha_bwd/source_3.cu", +] diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/dispatch_head_dim.h b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/dispatch_head_dim.h new file mode 100644 index 00000000..c7a5c533 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/dispatch_head_dim.h @@ -0,0 +1,93 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_blackwell { +#define DISPATCH_BLACKWELL_FMHA_FORWARD_float16(dim, q_tile_size, kv_tile_size, persistent, ...) \ + [&] { \ + if (dim <= 32) { \ + DISPATCH_BLACKWELL_FMHA_FORWARD_float16_headdim32(q_tile_size, kv_tile_size, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 64) { \ + DISPATCH_BLACKWELL_FMHA_FORWARD_float16_headdim64(q_tile_size, kv_tile_size, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 128) { \ + DISPATCH_BLACKWELL_FMHA_FORWARD_float16_headdim128(q_tile_size, kv_tile_size, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FMHA forward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for float16."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FMHA_FORWARD_bfloat16(dim, q_tile_size, kv_tile_size, persistent, ...) \ + [&] { \ + if (dim <= 32) { \ + DISPATCH_BLACKWELL_FMHA_FORWARD_bfloat16_headdim32(q_tile_size, kv_tile_size, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 64) { \ + DISPATCH_BLACKWELL_FMHA_FORWARD_bfloat16_headdim64(q_tile_size, kv_tile_size, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 128) { \ + DISPATCH_BLACKWELL_FMHA_FORWARD_bfloat16_headdim128(q_tile_size, kv_tile_size, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FMHA forward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for bfloat16."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FMHA_FORWARD_e4m3(dim, q_tile_size, kv_tile_size, persistent, ...) \ + [&] { \ + if (dim <= 32) { \ + DISPATCH_BLACKWELL_FMHA_FORWARD_e4m3_headdim32(q_tile_size, kv_tile_size, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 64) { \ + DISPATCH_BLACKWELL_FMHA_FORWARD_e4m3_headdim64(q_tile_size, kv_tile_size, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 128) { \ + DISPATCH_BLACKWELL_FMHA_FORWARD_e4m3_headdim128(q_tile_size, kv_tile_size, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FMHA forward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for e4m3."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FMHA_FORWARD_e5m2(dim, q_tile_size, kv_tile_size, persistent, ...) \ + [&] { \ + if (dim <= 32) { \ + DISPATCH_BLACKWELL_FMHA_FORWARD_e5m2_headdim32(q_tile_size, kv_tile_size, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 64) { \ + DISPATCH_BLACKWELL_FMHA_FORWARD_e5m2_headdim64(q_tile_size, kv_tile_size, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 128) { \ + DISPATCH_BLACKWELL_FMHA_FORWARD_e5m2_headdim128(q_tile_size, kv_tile_size, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FMHA forward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for e5m2."); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fmha_blackwell +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/dispatch_tile_size.h b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/dispatch_tile_size.h new file mode 100644 index 00000000..0fd5cc8c --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/dispatch_tile_size.h @@ -0,0 +1,209 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_blackwell { +#define DISPATCH_BLACKWELL_FMHA_FORWARD_float16_headdim32(q_tile_size, kv_tile_size, persistent, ...) \ + [&] { \ + if (q_tile_size == 256 && \ +kv_tile_size == 128) { \ + if (persistent) { \ + natten::cuda::fmha_blackwell::blackwell_fmha_float16_256x128x32_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fmha_blackwell::blackwell_fmha_float16_256x128x32(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FMHA forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + "."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FMHA_FORWARD_float16_headdim64(q_tile_size, kv_tile_size, persistent, ...) \ + [&] { \ + if (q_tile_size == 256 && \ +kv_tile_size == 128) { \ + if (persistent) { \ + natten::cuda::fmha_blackwell::blackwell_fmha_float16_256x128x64_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fmha_blackwell::blackwell_fmha_float16_256x128x64(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FMHA forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + "."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FMHA_FORWARD_float16_headdim128(q_tile_size, kv_tile_size, persistent, ...) \ + [&] { \ + if (q_tile_size == 256 && \ +kv_tile_size == 128) { \ + if (persistent) { \ + natten::cuda::fmha_blackwell::blackwell_fmha_float16_256x128x128_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fmha_blackwell::blackwell_fmha_float16_256x128x128(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FMHA forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + "."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FMHA_FORWARD_bfloat16_headdim32(q_tile_size, kv_tile_size, persistent, ...) \ + [&] { \ + if (q_tile_size == 256 && \ +kv_tile_size == 128) { \ + if (persistent) { \ + natten::cuda::fmha_blackwell::blackwell_fmha_bfloat16_256x128x32_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fmha_blackwell::blackwell_fmha_bfloat16_256x128x32(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FMHA forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + "."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FMHA_FORWARD_bfloat16_headdim64(q_tile_size, kv_tile_size, persistent, ...) \ + [&] { \ + if (q_tile_size == 256 && \ +kv_tile_size == 128) { \ + if (persistent) { \ + natten::cuda::fmha_blackwell::blackwell_fmha_bfloat16_256x128x64_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fmha_blackwell::blackwell_fmha_bfloat16_256x128x64(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FMHA forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + "."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FMHA_FORWARD_bfloat16_headdim128(q_tile_size, kv_tile_size, persistent, ...) \ + [&] { \ + if (q_tile_size == 256 && \ +kv_tile_size == 128) { \ + if (persistent) { \ + natten::cuda::fmha_blackwell::blackwell_fmha_bfloat16_256x128x128_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fmha_blackwell::blackwell_fmha_bfloat16_256x128x128(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FMHA forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + "."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FMHA_FORWARD_e4m3_headdim32(q_tile_size, kv_tile_size, persistent, ...) \ + [&] { \ + if (q_tile_size == 256 && \ +kv_tile_size == 128) { \ + if (persistent) { \ + natten::cuda::fmha_blackwell::blackwell_fmha_e4m3_256x128x32_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fmha_blackwell::blackwell_fmha_e4m3_256x128x32(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FMHA forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 32): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + "."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FMHA_FORWARD_e4m3_headdim64(q_tile_size, kv_tile_size, persistent, ...) \ + [&] { \ + if (q_tile_size == 256 && \ +kv_tile_size == 128) { \ + if (persistent) { \ + natten::cuda::fmha_blackwell::blackwell_fmha_e4m3_256x128x64_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fmha_blackwell::blackwell_fmha_e4m3_256x128x64(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FMHA forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 64): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + "."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FMHA_FORWARD_e4m3_headdim128(q_tile_size, kv_tile_size, persistent, ...) \ + [&] { \ + if (q_tile_size == 256 && \ +kv_tile_size == 128) { \ + if (persistent) { \ + natten::cuda::fmha_blackwell::blackwell_fmha_e4m3_256x128x128_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fmha_blackwell::blackwell_fmha_e4m3_256x128x128(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FMHA forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 128): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + "."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FMHA_FORWARD_e5m2_headdim32(q_tile_size, kv_tile_size, persistent, ...) \ + [&] { \ + if (q_tile_size == 256 && \ +kv_tile_size == 128) { \ + if (persistent) { \ + natten::cuda::fmha_blackwell::blackwell_fmha_e5m2_256x128x32_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fmha_blackwell::blackwell_fmha_e5m2_256x128x32(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FMHA forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 32): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + "."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FMHA_FORWARD_e5m2_headdim64(q_tile_size, kv_tile_size, persistent, ...) \ + [&] { \ + if (q_tile_size == 256 && \ +kv_tile_size == 128) { \ + if (persistent) { \ + natten::cuda::fmha_blackwell::blackwell_fmha_e5m2_256x128x64_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fmha_blackwell::blackwell_fmha_e5m2_256x128x64(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FMHA forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 64): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + "."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FMHA_FORWARD_e5m2_headdim128(q_tile_size, kv_tile_size, persistent, ...) \ + [&] { \ + if (q_tile_size == 256 && \ +kv_tile_size == 128) { \ + if (persistent) { \ + natten::cuda::fmha_blackwell::blackwell_fmha_e5m2_256x128x128_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fmha_blackwell::blackwell_fmha_e5m2_256x128x128(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FMHA forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 128): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + "."); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fmha_blackwell +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/interface.h b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/interface.h new file mode 100644 index 00000000..17357193 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/interface.h @@ -0,0 +1,48 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_blackwell { +#define DISPATCH_BLACKWELL_FMHA_FORWARD(dtype, dim, q_tile_size, kv_tile_size, persistent, ...) \ + [&] { \ + if (dtype == torch::kFloat16) { \ + DISPATCH_BLACKWELL_FMHA_FORWARD_float16(dim, q_tile_size, kv_tile_size, persistent, __VA_ARGS__); \ + } \ + else if (dtype == torch::kBFloat16) { \ + DISPATCH_BLACKWELL_FMHA_FORWARD_bfloat16(dim, q_tile_size, kv_tile_size, persistent, __VA_ARGS__); \ + } \ + else if (dtype == c10::ScalarType::Float8_e4m3fn) { \ + DISPATCH_BLACKWELL_FMHA_FORWARD_e4m3(dim, q_tile_size, kv_tile_size, persistent, __VA_ARGS__); \ + } \ + else if (dtype == c10::ScalarType::Float8_e5m2) { \ + DISPATCH_BLACKWELL_FMHA_FORWARD_e5m2(dim, q_tile_size, kv_tile_size, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FMHA forward kernel dispatch failed! It does not support dtype " + std::string(c10::toString(dtype)) + "."); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fmha_blackwell +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/kernels.h b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/kernels.h new file mode 100644 index 00000000..e775596f --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/kernels.h @@ -0,0 +1,628 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_blackwell { + +void blackwell_fmha_float16_256x128x32( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fmha_float16_256x128x32_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fmha_float16_256x128x64( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fmha_float16_256x128x64_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fmha_float16_256x128x128( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fmha_float16_256x128x128_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fmha_bfloat16_256x128x32( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fmha_bfloat16_256x128x32_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fmha_bfloat16_256x128x64( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fmha_bfloat16_256x128x64_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fmha_bfloat16_256x128x128( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fmha_bfloat16_256x128x128_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fmha_e4m3_256x128x32( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fmha_e4m3_256x128x32_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fmha_e4m3_256x128x64( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fmha_e4m3_256x128x64_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fmha_e4m3_256x128x128( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fmha_e4m3_256x128x128_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fmha_e5m2_256x128x32( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fmha_e5m2_256x128x32_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fmha_e5m2_256x128x64( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fmha_e5m2_256x128x64_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fmha_e5m2_256x128x128( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fmha_e5m2_256x128x128_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + + +} // namespace natten +} // namespace cuda +} // namespace fmha_blackwell +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/dispatch_head_dim.h b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/dispatch_head_dim.h new file mode 100644 index 00000000..66100f8a --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/dispatch_head_dim.h @@ -0,0 +1,61 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_blackwell { +#define DISPATCH_BLACKWELL_FMHA_BACKWARD_float16(dim, q_tile_size, kv_tile_size, ...) \ + [&] { \ + if (dim <= 32) { \ + DISPATCH_BLACKWELL_FMHA_BACKWARD_float16_headdim32(q_tile_size, kv_tile_size, __VA_ARGS__); \ + } \ + else if (dim <= 64) { \ + DISPATCH_BLACKWELL_FMHA_BACKWARD_float16_headdim64(q_tile_size, kv_tile_size, __VA_ARGS__); \ + } \ + else if (dim <= 128) { \ + DISPATCH_BLACKWELL_FMHA_BACKWARD_float16_headdim128(q_tile_size, kv_tile_size, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FMHA backward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for float16."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FMHA_BACKWARD_bfloat16(dim, q_tile_size, kv_tile_size, ...) \ + [&] { \ + if (dim <= 32) { \ + DISPATCH_BLACKWELL_FMHA_BACKWARD_bfloat16_headdim32(q_tile_size, kv_tile_size, __VA_ARGS__); \ + } \ + else if (dim <= 64) { \ + DISPATCH_BLACKWELL_FMHA_BACKWARD_bfloat16_headdim64(q_tile_size, kv_tile_size, __VA_ARGS__); \ + } \ + else if (dim <= 128) { \ + DISPATCH_BLACKWELL_FMHA_BACKWARD_bfloat16_headdim128(q_tile_size, kv_tile_size, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FMHA backward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for bfloat16."); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fmha_blackwell +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/dispatch_tile_size.h b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/dispatch_tile_size.h new file mode 100644 index 00000000..3b659ee8 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/dispatch_tile_size.h @@ -0,0 +1,95 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_blackwell { +#define DISPATCH_BLACKWELL_FMHA_BACKWARD_float16_headdim32(q_tile_size, kv_tile_size, ...) \ + [&] { \ + if (q_tile_size == 128 && \ +kv_tile_size == 128) { \ + natten::cuda::fmha_blackwell::blackwell_fmha_backward_float16_128x128x32(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FMHA backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + "."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FMHA_BACKWARD_float16_headdim64(q_tile_size, kv_tile_size, ...) \ + [&] { \ + if (q_tile_size == 128 && \ +kv_tile_size == 128) { \ + natten::cuda::fmha_blackwell::blackwell_fmha_backward_float16_128x128x64(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FMHA backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + "."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FMHA_BACKWARD_float16_headdim128(q_tile_size, kv_tile_size, ...) \ + [&] { \ + if (q_tile_size == 128 && \ +kv_tile_size == 128) { \ + natten::cuda::fmha_blackwell::blackwell_fmha_backward_float16_128x128x128(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FMHA backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + "."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FMHA_BACKWARD_bfloat16_headdim32(q_tile_size, kv_tile_size, ...) \ + [&] { \ + if (q_tile_size == 128 && \ +kv_tile_size == 128) { \ + natten::cuda::fmha_blackwell::blackwell_fmha_backward_bfloat16_128x128x32(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FMHA backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + "."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FMHA_BACKWARD_bfloat16_headdim64(q_tile_size, kv_tile_size, ...) \ + [&] { \ + if (q_tile_size == 128 && \ +kv_tile_size == 128) { \ + natten::cuda::fmha_blackwell::blackwell_fmha_backward_bfloat16_128x128x64(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FMHA backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + "."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FMHA_BACKWARD_bfloat16_headdim128(q_tile_size, kv_tile_size, ...) \ + [&] { \ + if (q_tile_size == 128 && \ +kv_tile_size == 128) { \ + natten::cuda::fmha_blackwell::blackwell_fmha_backward_bfloat16_128x128x128(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FMHA backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + "."); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fmha_blackwell +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/interface.h b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/interface.h new file mode 100644 index 00000000..103277bc --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/interface.h @@ -0,0 +1,42 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_blackwell { +#define DISPATCH_BLACKWELL_FMHA_BACKWARD(dtype, dim, q_tile_size, kv_tile_size, ...) \ + [&] { \ + if (dtype == torch::kFloat16) { \ + DISPATCH_BLACKWELL_FMHA_BACKWARD_float16(dim, q_tile_size, kv_tile_size, __VA_ARGS__); \ + } \ + else if (dtype == torch::kBFloat16) { \ + DISPATCH_BLACKWELL_FMHA_BACKWARD_bfloat16(dim, q_tile_size, kv_tile_size, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FMHA backward kernel dispatch failed! It does not support dtype " + std::string(c10::toString(dtype)) + "."); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fmha_blackwell +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/kernels.h b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/kernels.h new file mode 100644 index 00000000..3e030b82 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/kernels.h @@ -0,0 +1,208 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_blackwell { + +void blackwell_fmha_backward_float16_128x128x32( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + bool deterministic, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fmha_backward_float16_128x128x64( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + bool deterministic, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fmha_backward_float16_128x128x128( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + bool deterministic, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fmha_backward_bfloat16_128x128x32( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + bool deterministic, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fmha_backward_bfloat16_128x128x64( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + bool deterministic, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fmha_backward_bfloat16_128x128x128( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + bool deterministic, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + + +} // namespace natten +} // namespace cuda +} // namespace fmha_blackwell +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_cm.h b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_cm.h new file mode 100644 index 00000000..a28e90ed --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_cm.h @@ -0,0 +1,785 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_float16_headdim32(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_float16_headdim32_causal0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_float16_headdim32_causal1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_float16_headdim64(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_float16_headdim64_causal0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_float16_headdim64_causal1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_float16_headdim128(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_float16_headdim128_causal0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_float16_headdim128_causal1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_bfloat16_headdim32(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_bfloat16_headdim32_causal0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_bfloat16_headdim32_causal1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_bfloat16_headdim64(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_bfloat16_headdim64_causal0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_bfloat16_headdim64_causal1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_bfloat16_headdim128(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_bfloat16_headdim128_causal0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_bfloat16_headdim128_causal1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_e4m3_headdim32(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_e4m3_headdim32_causal0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_e4m3_headdim32_causal1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! Causal mask dispatcher (e4m3, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_e4m3_headdim64(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_e4m3_headdim64_causal0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_e4m3_headdim64_causal1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! Causal mask dispatcher (e4m3, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_e4m3_headdim128(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_e4m3_headdim128_causal0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_e4m3_headdim128_causal1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! Causal mask dispatcher (e4m3, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_e5m2_headdim32(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_e5m2_headdim32_causal0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_e5m2_headdim32_causal1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! Causal mask dispatcher (e5m2, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_e5m2_headdim64(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_e5m2_headdim64_causal0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_e5m2_headdim64_causal1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! Causal mask dispatcher (e5m2, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_e5m2_headdim128(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_e5m2_headdim128_causal0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_e5m2_headdim128_causal1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! Causal mask dispatcher (e5m2, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim32(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim32_causal0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim32_causal0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim32_causal1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim32_causal1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim64(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim64_causal0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim64_causal0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim64_causal1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim64_causal1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim128(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim128_causal0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim128_causal0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim128_causal1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim128_causal1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim32(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim32_causal0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim32_causal0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim32_causal1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim32_causal1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim64(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim64_causal0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim64_causal0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim64_causal1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim64_causal1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim128(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim128_causal0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim128_causal0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim128_causal1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim128_causal1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim32(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim32_causal0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim32_causal0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim32_causal1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim32_causal1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! Causal mask dispatcher (e4m3, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim64(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim64_causal0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim64_causal0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim64_causal1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim64_causal1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! Causal mask dispatcher (e4m3, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim128(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim128_causal0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim128_causal0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim128_causal1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim128_causal1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! Causal mask dispatcher (e4m3, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim32(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim32_causal0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim32_causal0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim32_causal1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim32_causal1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! Causal mask dispatcher (e5m2, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim64(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim64_causal0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim64_causal0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim64_causal1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim64_causal1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! Causal mask dispatcher (e5m2, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim128(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim128_causal0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim128_causal0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim128_causal1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim128_causal1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! Causal mask dispatcher (e5m2, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim32(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim32_causal0x0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim32_causal0x0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim32_causal0x1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim32_causal0x1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim32_causal1x0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim32_causal1x0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim32_causal1x1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim32_causal1x1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim64(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim64_causal0x0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim64_causal0x0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim64_causal0x1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim64_causal0x1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim64_causal1x0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim64_causal1x0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim64_causal1x1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim64_causal1x1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim128(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim128_causal0x0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim128_causal0x0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim128_causal0x1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim128_causal0x1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim128_causal1x0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim128_causal1x0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim128_causal1x1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim128_causal1x1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim32(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim32_causal0x0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim32_causal0x0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim32_causal0x1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim32_causal0x1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim32_causal1x0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim32_causal1x0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim32_causal1x1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim32_causal1x1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim64(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim64_causal0x0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim64_causal0x0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim64_causal0x1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim64_causal0x1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim64_causal1x0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim64_causal1x0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim64_causal1x1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim64_causal1x1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim128(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim128_causal0x0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim128_causal0x0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim128_causal0x1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim128_causal0x1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim128_causal1x0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim128_causal1x0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim128_causal1x1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim128_causal1x1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim32(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim32_causal0x0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim32_causal0x0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim32_causal0x1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim32_causal0x1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim32_causal1x0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim32_causal1x0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim32_causal1x1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim32_causal1x1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! Causal mask dispatcher (e4m3, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim64(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim64_causal0x0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim64_causal0x0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim64_causal0x1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim64_causal0x1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim64_causal1x0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim64_causal1x0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim64_causal1x1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim64_causal1x1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! Causal mask dispatcher (e4m3, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim128(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim128_causal0x0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim128_causal0x0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim128_causal0x1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim128_causal0x1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim128_causal1x0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim128_causal1x0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim128_causal1x1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim128_causal1x1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! Causal mask dispatcher (e4m3, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim32(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim32_causal0x0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim32_causal0x0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim32_causal0x1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim32_causal0x1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim32_causal1x0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim32_causal1x0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim32_causal1x1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim32_causal1x1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! Causal mask dispatcher (e5m2, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim64(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim64_causal0x0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim64_causal0x0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim64_causal0x1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim64_causal0x1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim64_causal1x0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim64_causal1x0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim64_causal1x1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim64_causal1x1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! Causal mask dispatcher (e5m2, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim128(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim128_causal0x0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim128_causal0x0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim128_causal0x1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim128_causal0x1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim128_causal1x0x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim128_causal1x0x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim128_causal1x1x0(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim128_causal1x1x1(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! Causal mask dispatcher (e5m2, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fna_blackwell +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_dtype.h b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_dtype.h new file mode 100644 index 00000000..2e3b6392 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_dtype.h @@ -0,0 +1,86 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (dtype == torch::kFloat16) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dtype == torch::kBFloat16) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dtype == c10::ScalarType::Float8_e4m3fn) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_e4m3(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dtype == c10::ScalarType::Float8_e5m2) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_e5m2(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! It does not support dtype " + std::string(c10::toString(dtype)) + "."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (dtype == torch::kFloat16) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dtype == torch::kBFloat16) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dtype == c10::ScalarType::Float8_e4m3fn) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dtype == c10::ScalarType::Float8_e5m2) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It does not support dtype " + std::string(c10::toString(dtype)) + "."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (dtype == torch::kFloat16) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dtype == torch::kBFloat16) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dtype == c10::ScalarType::Float8_e4m3fn) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dtype == c10::ScalarType::Float8_e5m2) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It does not support dtype " + std::string(c10::toString(dtype)) + "."); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fna_blackwell +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_head_dim.h b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_head_dim.h new file mode 100644 index 00000000..760a72f6 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_head_dim.h @@ -0,0 +1,221 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (dim <= 32) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_float16_headdim32(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 64) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_float16_headdim64(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 128) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_float16_headdim128(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA forward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for float16."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (dim <= 32) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_bfloat16_headdim32(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 64) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_bfloat16_headdim64(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 128) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_bfloat16_headdim128(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA forward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for bfloat16."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_e4m3(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (dim <= 32) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_e4m3_headdim32(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 64) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_e4m3_headdim64(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 128) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_e4m3_headdim128(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA forward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for e4m3."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_e5m2(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (dim <= 32) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_e5m2_headdim32(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 64) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_e5m2_headdim64(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 128) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D_e5m2_headdim128(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA forward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for e5m2."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (dim <= 32) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim32(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 64) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim64(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 128) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim128(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA forward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for float16."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (dim <= 32) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim32(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 64) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim64(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 128) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim128(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA forward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for bfloat16."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (dim <= 32) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim32(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 64) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim64(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 128) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim128(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA forward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for e4m3."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (dim <= 32) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim32(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 64) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim64(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 128) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim128(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA forward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for e5m2."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (dim <= 32) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim32(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 64) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim64(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 128) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim128(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA forward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for float16."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (dim <= 32) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim32(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 64) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim64(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 128) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim128(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA forward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for bfloat16."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (dim <= 32) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim32(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 64) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim64(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 128) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim128(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA forward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for e4m3."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (dim <= 32) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim32(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 64) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim64(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if (dim <= 128) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim128(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA forward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for e5m2."); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fna_blackwell +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_tile_shape.h b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_tile_shape.h new file mode 100644 index 00000000..8faccc9e --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_tile_shape.h @@ -0,0 +1,8309 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_float16_headdim32_causal0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 256 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_float16_256x128x32_Q256_KV128_causal0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna1d_float16_256x128x32_Q256_KV128_causal0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_float16_headdim32_causal1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 256 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_float16_256x128x32_Q256_KV128_causal1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna1d_float16_256x128x32_Q256_KV128_causal1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_float16_headdim64_causal0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 256 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_float16_256x128x64_Q256_KV128_causal0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna1d_float16_256x128x64_Q256_KV128_causal0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_float16_headdim64_causal1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 256 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_float16_256x128x64_Q256_KV128_causal1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna1d_float16_256x128x64_Q256_KV128_causal1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_float16_headdim128_causal0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 256 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_float16_256x128x128_Q256_KV128_causal0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna1d_float16_256x128x128_Q256_KV128_causal0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_float16_headdim128_causal1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 256 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_float16_256x128x128_Q256_KV128_causal1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna1d_float16_256x128x128_Q256_KV128_causal1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_bfloat16_headdim32_causal0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 256 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_bfloat16_256x128x32_Q256_KV128_causal0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna1d_bfloat16_256x128x32_Q256_KV128_causal0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_bfloat16_headdim32_causal1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 256 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_bfloat16_256x128x32_Q256_KV128_causal1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna1d_bfloat16_256x128x32_Q256_KV128_causal1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_bfloat16_headdim64_causal0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 256 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_bfloat16_256x128x64_Q256_KV128_causal0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna1d_bfloat16_256x128x64_Q256_KV128_causal0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_bfloat16_headdim64_causal1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 256 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_bfloat16_256x128x64_Q256_KV128_causal1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna1d_bfloat16_256x128x64_Q256_KV128_causal1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_bfloat16_headdim128_causal0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 256 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_bfloat16_256x128x128_Q256_KV128_causal0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna1d_bfloat16_256x128x128_Q256_KV128_causal0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_bfloat16_headdim128_causal1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 256 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_bfloat16_256x128x128_Q256_KV128_causal1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna1d_bfloat16_256x128x128_Q256_KV128_causal1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_e4m3_headdim32_causal0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 256 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_e4m3_256x128x32_Q256_KV128_causal0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna1d_e4m3_256x128x32_Q256_KV128_causal0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_e4m3_headdim32_causal1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 256 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_e4m3_256x128x32_Q256_KV128_causal1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna1d_e4m3_256x128x32_Q256_KV128_causal1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_e4m3_headdim64_causal0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 256 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_e4m3_256x128x64_Q256_KV128_causal0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna1d_e4m3_256x128x64_Q256_KV128_causal0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_e4m3_headdim64_causal1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 256 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_e4m3_256x128x64_Q256_KV128_causal1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna1d_e4m3_256x128x64_Q256_KV128_causal1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_e4m3_headdim128_causal0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 256 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_e4m3_256x128x128_Q256_KV128_causal0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna1d_e4m3_256x128x128_Q256_KV128_causal0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_e4m3_headdim128_causal1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 256 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_e4m3_256x128x128_Q256_KV128_causal1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna1d_e4m3_256x128x128_Q256_KV128_causal1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_e5m2_headdim32_causal0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 256 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_e5m2_256x128x32_Q256_KV128_causal0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna1d_e5m2_256x128x32_Q256_KV128_causal0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_e5m2_headdim32_causal1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 256 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_e5m2_256x128x32_Q256_KV128_causal1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna1d_e5m2_256x128x32_Q256_KV128_causal1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_e5m2_headdim64_causal0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 256 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_e5m2_256x128x64_Q256_KV128_causal0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna1d_e5m2_256x128x64_Q256_KV128_causal0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_e5m2_headdim64_causal1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 256 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_e5m2_256x128x64_Q256_KV128_causal1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna1d_e5m2_256x128x64_Q256_KV128_causal1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_e5m2_headdim128_causal0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 256 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_e5m2_256x128x128_Q256_KV128_causal0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna1d_e5m2_256x128x128_Q256_KV128_causal0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_1D_e5m2_headdim128_causal1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 256 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_e5m2_256x128x128_Q256_KV128_causal1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna1d_e5m2_256x128x128_Q256_KV128_causal1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim32_causal0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q16x16_KV16x8_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q16x16_KV16x8_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q16x16_KV8x16_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q16x16_KV8x16_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q8x32_KV8x16_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q8x32_KV8x16_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q8x32_KV4x32_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q8x32_KV4x32_causal0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim32_causal0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q16x16_KV16x8_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q16x16_KV16x8_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q16x16_KV8x16_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q16x16_KV8x16_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q8x32_KV8x16_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q8x32_KV8x16_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q8x32_KV4x32_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q8x32_KV4x32_causal0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim32_causal1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q16x16_KV16x8_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q16x16_KV16x8_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q16x16_KV8x16_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q16x16_KV8x16_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q8x32_KV8x16_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q8x32_KV8x16_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q8x32_KV4x32_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q8x32_KV4x32_causal1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim32_causal1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q16x16_KV16x8_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q16x16_KV16x8_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q16x16_KV8x16_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q16x16_KV8x16_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q8x32_KV8x16_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q8x32_KV8x16_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q8x32_KV4x32_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x32_Q8x32_KV4x32_causal1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim64_causal0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q16x16_KV16x8_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q16x16_KV16x8_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q16x16_KV8x16_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q16x16_KV8x16_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q8x32_KV8x16_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q8x32_KV8x16_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q8x32_KV4x32_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q8x32_KV4x32_causal0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim64_causal0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q16x16_KV16x8_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q16x16_KV16x8_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q16x16_KV8x16_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q16x16_KV8x16_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q8x32_KV8x16_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q8x32_KV8x16_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q8x32_KV4x32_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q8x32_KV4x32_causal0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim64_causal1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q16x16_KV16x8_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q16x16_KV16x8_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q16x16_KV8x16_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q16x16_KV8x16_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q8x32_KV8x16_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q8x32_KV8x16_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q8x32_KV4x32_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q8x32_KV4x32_causal1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim64_causal1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q16x16_KV16x8_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q16x16_KV16x8_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q16x16_KV8x16_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q16x16_KV8x16_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q8x32_KV8x16_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q8x32_KV8x16_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q8x32_KV4x32_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x64_Q8x32_KV4x32_causal1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim128_causal0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q16x16_KV16x8_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q16x16_KV16x8_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q16x16_KV8x16_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q16x16_KV8x16_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q8x32_KV8x16_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q8x32_KV8x16_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q8x32_KV4x32_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q8x32_KV4x32_causal0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim128_causal0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q16x16_KV16x8_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q16x16_KV16x8_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q16x16_KV8x16_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q16x16_KV8x16_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q8x32_KV8x16_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q8x32_KV8x16_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q8x32_KV4x32_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q8x32_KV4x32_causal0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim128_causal1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q16x16_KV16x8_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q16x16_KV16x8_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q16x16_KV8x16_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q16x16_KV8x16_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q8x32_KV8x16_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q8x32_KV8x16_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q8x32_KV4x32_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q8x32_KV4x32_causal1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16_headdim128_causal1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q16x16_KV16x8_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q16x16_KV16x8_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q16x16_KV8x16_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q16x16_KV8x16_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q8x32_KV8x16_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q8x32_KV8x16_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q8x32_KV4x32_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_float16_256x128x128_Q8x32_KV4x32_causal1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim32_causal0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV16x8_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV16x8_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV8x16_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV8x16_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV8x16_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV8x16_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV4x32_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV4x32_causal0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim32_causal0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV16x8_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV16x8_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV8x16_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV8x16_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV8x16_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV8x16_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV4x32_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV4x32_causal0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim32_causal1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV16x8_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV16x8_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV8x16_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV8x16_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV8x16_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV8x16_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV4x32_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV4x32_causal1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim32_causal1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV16x8_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV16x8_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV8x16_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV8x16_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV8x16_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV8x16_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV4x32_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV4x32_causal1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim64_causal0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV16x8_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV16x8_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV8x16_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV8x16_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV8x16_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV8x16_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV4x32_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV4x32_causal0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim64_causal0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV16x8_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV16x8_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV8x16_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV8x16_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV8x16_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV8x16_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV4x32_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV4x32_causal0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim64_causal1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV16x8_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV16x8_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV8x16_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV8x16_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV8x16_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV8x16_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV4x32_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV4x32_causal1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim64_causal1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV16x8_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV16x8_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV8x16_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV8x16_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV8x16_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV8x16_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV4x32_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV4x32_causal1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim128_causal0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV16x8_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV16x8_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV8x16_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV8x16_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV8x16_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV8x16_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV4x32_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV4x32_causal0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim128_causal0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV16x8_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV16x8_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV8x16_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV8x16_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV8x16_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV8x16_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV4x32_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV4x32_causal0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim128_causal1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV16x8_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV16x8_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV8x16_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV8x16_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV8x16_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV8x16_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV4x32_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV4x32_causal1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16_headdim128_causal1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV16x8_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV16x8_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV8x16_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV8x16_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV8x16_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV8x16_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV4x32_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV4x32_causal1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim32_causal0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q16x16_KV16x8_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q16x16_KV16x8_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q16x16_KV8x16_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q16x16_KV8x16_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q8x32_KV8x16_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q8x32_KV8x16_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q8x32_KV4x32_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q8x32_KV4x32_causal0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim32_causal0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q16x16_KV16x8_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q16x16_KV16x8_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q16x16_KV8x16_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q16x16_KV8x16_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q8x32_KV8x16_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q8x32_KV8x16_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q8x32_KV4x32_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q8x32_KV4x32_causal0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim32_causal1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q16x16_KV16x8_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q16x16_KV16x8_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q16x16_KV8x16_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q16x16_KV8x16_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q8x32_KV8x16_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q8x32_KV8x16_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q8x32_KV4x32_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q8x32_KV4x32_causal1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim32_causal1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q16x16_KV16x8_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q16x16_KV16x8_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q16x16_KV8x16_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q16x16_KV8x16_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q8x32_KV8x16_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q8x32_KV8x16_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q8x32_KV4x32_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x32_Q8x32_KV4x32_causal1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim64_causal0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q16x16_KV16x8_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q16x16_KV16x8_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q16x16_KV8x16_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q16x16_KV8x16_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q8x32_KV8x16_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q8x32_KV8x16_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q8x32_KV4x32_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q8x32_KV4x32_causal0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim64_causal0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q16x16_KV16x8_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q16x16_KV16x8_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q16x16_KV8x16_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q16x16_KV8x16_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q8x32_KV8x16_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q8x32_KV8x16_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q8x32_KV4x32_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q8x32_KV4x32_causal0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim64_causal1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q16x16_KV16x8_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q16x16_KV16x8_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q16x16_KV8x16_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q16x16_KV8x16_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q8x32_KV8x16_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q8x32_KV8x16_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q8x32_KV4x32_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q8x32_KV4x32_causal1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim64_causal1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q16x16_KV16x8_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q16x16_KV16x8_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q16x16_KV8x16_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q16x16_KV8x16_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q8x32_KV8x16_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q8x32_KV8x16_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q8x32_KV4x32_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x64_Q8x32_KV4x32_causal1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim128_causal0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q16x16_KV16x8_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q16x16_KV16x8_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q16x16_KV8x16_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q16x16_KV8x16_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q8x32_KV8x16_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q8x32_KV8x16_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q8x32_KV4x32_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q8x32_KV4x32_causal0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim128_causal0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q16x16_KV16x8_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q16x16_KV16x8_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q16x16_KV8x16_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q16x16_KV8x16_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q8x32_KV8x16_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q8x32_KV8x16_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q8x32_KV4x32_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q8x32_KV4x32_causal0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim128_causal1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q16x16_KV16x8_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q16x16_KV16x8_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q16x16_KV8x16_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q16x16_KV8x16_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q8x32_KV8x16_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q8x32_KV8x16_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q8x32_KV4x32_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q8x32_KV4x32_causal1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e4m3_headdim128_causal1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q16x16_KV16x8_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q16x16_KV16x8_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q16x16_KV8x16_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q16x16_KV8x16_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q8x32_KV8x16_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q8x32_KV8x16_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q8x32_KV4x32_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e4m3_256x128x128_Q8x32_KV4x32_causal1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim32_causal0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q16x16_KV16x8_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q16x16_KV16x8_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q16x16_KV8x16_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q16x16_KV8x16_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q8x32_KV8x16_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q8x32_KV8x16_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q8x32_KV4x32_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q8x32_KV4x32_causal0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim32_causal0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q16x16_KV16x8_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q16x16_KV16x8_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q16x16_KV8x16_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q16x16_KV8x16_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q8x32_KV8x16_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q8x32_KV8x16_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q8x32_KV4x32_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q8x32_KV4x32_causal0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim32_causal1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q16x16_KV16x8_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q16x16_KV16x8_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q16x16_KV8x16_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q16x16_KV8x16_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q8x32_KV8x16_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q8x32_KV8x16_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q8x32_KV4x32_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q8x32_KV4x32_causal1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim32_causal1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q16x16_KV16x8_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q16x16_KV16x8_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q16x16_KV8x16_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q16x16_KV8x16_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q8x32_KV8x16_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q8x32_KV8x16_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q8x32_KV4x32_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x32_Q8x32_KV4x32_causal1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim64_causal0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q16x16_KV16x8_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q16x16_KV16x8_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q16x16_KV8x16_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q16x16_KV8x16_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q8x32_KV8x16_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q8x32_KV8x16_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q8x32_KV4x32_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q8x32_KV4x32_causal0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim64_causal0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q16x16_KV16x8_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q16x16_KV16x8_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q16x16_KV8x16_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q16x16_KV8x16_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q8x32_KV8x16_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q8x32_KV8x16_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q8x32_KV4x32_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q8x32_KV4x32_causal0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim64_causal1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q16x16_KV16x8_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q16x16_KV16x8_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q16x16_KV8x16_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q16x16_KV8x16_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q8x32_KV8x16_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q8x32_KV8x16_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q8x32_KV4x32_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q8x32_KV4x32_causal1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim64_causal1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q16x16_KV16x8_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q16x16_KV16x8_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q16x16_KV8x16_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q16x16_KV8x16_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q8x32_KV8x16_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q8x32_KV8x16_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q8x32_KV4x32_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x64_Q8x32_KV4x32_causal1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim128_causal0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q16x16_KV16x8_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q16x16_KV16x8_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q16x16_KV8x16_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q16x16_KV8x16_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q8x32_KV8x16_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q8x32_KV8x16_causal0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q8x32_KV4x32_causal0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q8x32_KV4x32_causal0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim128_causal0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q16x16_KV16x8_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q16x16_KV16x8_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q16x16_KV8x16_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q16x16_KV8x16_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q8x32_KV8x16_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q8x32_KV8x16_causal0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q8x32_KV4x32_causal0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q8x32_KV4x32_causal0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim128_causal1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q16x16_KV16x8_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q16x16_KV16x8_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q16x16_KV8x16_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q16x16_KV8x16_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q8x32_KV8x16_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q8x32_KV8x16_causal1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q8x32_KV4x32_causal1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q8x32_KV4x32_causal1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_2D_e5m2_headdim128_causal1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q16x16_KV16x8_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q16x16_KV16x8_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q16x16_KV8x16_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q16x16_KV8x16_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q8x32_KV8x16_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q8x32_KV8x16_causal1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 32 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 32) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q8x32_KV4x32_causal1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna2d_e5m2_256x128x128_Q8x32_KV4x32_causal1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim32_causal0x0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim32_causal0x0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim32_causal0x1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim32_causal0x1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim32_causal1x0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim32_causal1x0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim32_causal1x1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim32_causal1x1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim64_causal0x0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim64_causal0x0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim64_causal0x1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim64_causal0x1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim64_causal1x0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim64_causal1x0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim64_causal1x1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim64_causal1x1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim128_causal0x0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim128_causal0x0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim128_causal0x1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim128_causal0x1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim128_causal1x0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim128_causal1x0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim128_causal1x1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16_headdim128_causal1x1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim32_causal0x0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim32_causal0x0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim32_causal0x1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim32_causal0x1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim32_causal1x0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim32_causal1x0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim32_causal1x1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim32_causal1x1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim64_causal0x0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim64_causal0x0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim64_causal0x1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim64_causal0x1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim64_causal1x0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim64_causal1x0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim64_causal1x1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim64_causal1x1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim128_causal0x0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim128_causal0x0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim128_causal0x1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim128_causal0x1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim128_causal1x0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim128_causal1x0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim128_causal1x1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16_headdim128_causal1x1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim32_causal0x0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim32_causal0x0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim32_causal0x1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim32_causal0x1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim32_causal1x0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim32_causal1x0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim32_causal1x1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim32_causal1x1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim64_causal0x0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim64_causal0x0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim64_causal0x1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim64_causal0x1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim64_causal1x0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim64_causal1x0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim64_causal1x1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim64_causal1x1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim128_causal0x0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim128_causal0x0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim128_causal0x1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim128_causal0x1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim128_causal1x0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim128_causal1x0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim128_causal1x1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e4m3_headdim128_causal1x1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e4m3, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim32_causal0x0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim32_causal0x0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim32_causal0x1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim32_causal0x1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim32_causal1x0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim32_causal1x0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim32_causal1x1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim32_causal1x1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim64_causal0x0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim64_causal0x0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim64_causal0x1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim64_causal0x1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim64_causal1x0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim64_causal1x0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim64_causal1x1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim64_causal1x1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim128_causal0x0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal0x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim128_causal0x0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal0x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim128_causal0x1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal0x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim128_causal0x1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal0x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim128_causal1x0x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal1x0x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim128_causal1x0x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal1x0x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim128_causal1x1x0(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal1x1x0_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_FORWARD_3D_e5m2_headdim128_causal1x1x1(q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 16 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + if (persistent) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal1x1x1_persistent(__VA_ARGS__); \ + } else { \ + natten::cuda::fna_blackwell::blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ + } \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D forward kernel dispatch failed! It got invalid Q tile and KV tile combination (e5m2, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fna_blackwell +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna/interface.h b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna/interface.h new file mode 100644 index 00000000..de17cf0f --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna/interface.h @@ -0,0 +1,45 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { +#define DISPATCH_BLACKWELL_FNA_FORWARD(rank, dtype, dim, is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ + [&] { \ + if constexpr (rank == 1) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_1D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if constexpr (rank == 2) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_2D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else if constexpr (rank == 3) { \ + DISPATCH_BLACKWELL_FNA_FORWARD_3D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA forward kernel dispatch failed! It only supports NA1D, 2D, and 3D!"); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fna_blackwell +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna/kernels.h b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna/kernels.h new file mode 100644 index 00000000..8c5e4f5a --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna/kernels.h @@ -0,0 +1,44428 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + +void blackwell_fna1d_float16_256x128x32_Q256_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_float16_256x128x32_Q256_KV128_causal0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_float16_256x128x32_Q256_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_float16_256x128x32_Q256_KV128_causal1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_float16_256x128x64_Q256_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_float16_256x128x64_Q256_KV128_causal0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_float16_256x128x64_Q256_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_float16_256x128x64_Q256_KV128_causal1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_float16_256x128x128_Q256_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_float16_256x128x128_Q256_KV128_causal0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_float16_256x128x128_Q256_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_float16_256x128x128_Q256_KV128_causal1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_bfloat16_256x128x32_Q256_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_bfloat16_256x128x32_Q256_KV128_causal0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_bfloat16_256x128x32_Q256_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_bfloat16_256x128x32_Q256_KV128_causal1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_bfloat16_256x128x64_Q256_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_bfloat16_256x128x64_Q256_KV128_causal0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_bfloat16_256x128x64_Q256_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_bfloat16_256x128x64_Q256_KV128_causal1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_bfloat16_256x128x128_Q256_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_bfloat16_256x128x128_Q256_KV128_causal0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_bfloat16_256x128x128_Q256_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_bfloat16_256x128x128_Q256_KV128_causal1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_e4m3_256x128x32_Q256_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_e4m3_256x128x32_Q256_KV128_causal0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_e4m3_256x128x32_Q256_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_e4m3_256x128x32_Q256_KV128_causal1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_e4m3_256x128x64_Q256_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_e4m3_256x128x64_Q256_KV128_causal0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_e4m3_256x128x64_Q256_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_e4m3_256x128x64_Q256_KV128_causal1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_e4m3_256x128x128_Q256_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_e4m3_256x128x128_Q256_KV128_causal0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_e4m3_256x128x128_Q256_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_e4m3_256x128x128_Q256_KV128_causal1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_e5m2_256x128x32_Q256_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_e5m2_256x128x32_Q256_KV128_causal0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_e5m2_256x128x32_Q256_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_e5m2_256x128x32_Q256_KV128_causal1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_e5m2_256x128x64_Q256_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_e5m2_256x128x64_Q256_KV128_causal0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_e5m2_256x128x64_Q256_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_e5m2_256x128x64_Q256_KV128_causal1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_e5m2_256x128x128_Q256_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_e5m2_256x128x128_Q256_KV128_causal0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_e5m2_256x128x128_Q256_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_e5m2_256x128x128_Q256_KV128_causal1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV16x8_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV4x32_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV4x32_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV16x8_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV4x32_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV4x32_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV16x8_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV4x32_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV4x32_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV16x8_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV4x32_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV4x32_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV16x8_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV4x32_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV4x32_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV16x8_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV4x32_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV4x32_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV16x8_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV4x32_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV4x32_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV16x8_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV4x32_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV4x32_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV16x8_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV4x32_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV4x32_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV16x8_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV4x32_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV4x32_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV16x8_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV4x32_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV4x32_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV16x8_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV4x32_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV4x32_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV16x8_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV4x32_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV4x32_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV16x8_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV4x32_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV4x32_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV16x8_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV4x32_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV4x32_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV16x8_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV4x32_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV4x32_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV16x8_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV4x32_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV4x32_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV16x8_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV4x32_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV4x32_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV16x8_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV4x32_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV4x32_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV16x8_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV4x32_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV4x32_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV16x8_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV4x32_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV4x32_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV16x8_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV4x32_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV4x32_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV16x8_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV4x32_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV4x32_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV16x8_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV4x32_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV4x32_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV16x8_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV4x32_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV4x32_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV16x8_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV4x32_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV4x32_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV16x8_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV4x32_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV4x32_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV16x8_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV4x32_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV4x32_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV16x8_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV4x32_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV4x32_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV16x8_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV4x32_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV4x32_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV16x8_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV4x32_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV4x32_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV16x8_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV4x32_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV4x32_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV16x8_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV4x32_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV4x32_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV16x8_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV4x32_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV4x32_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV16x8_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV4x32_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV4x32_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV16x8_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV4x32_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV4x32_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV16x8_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV4x32_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV4x32_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV16x8_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV4x32_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV4x32_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV16x8_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV4x32_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV4x32_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV16x8_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV4x32_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV4x32_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV16x8_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV4x32_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV4x32_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV16x8_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV4x32_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV4x32_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV16x8_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV4x32_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV4x32_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV16x8_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV4x32_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV4x32_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV16x8_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV4x32_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV4x32_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV16x8_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV4x32_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV4x32_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV16x8_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV4x32_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV4x32_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV16x8_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV4x32_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV4x32_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + + +} // namespace natten +} // namespace cuda +} // namespace fna_blackwell +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_cm.h b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_cm.h new file mode 100644 index 00000000..e7b920da --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_cm.h @@ -0,0 +1,407 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { +#define DISPATCH_BLACKWELL_FNA_BACKWARD_1D_float16_headdim32(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_1D_float16_headdim32_causal0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_1D_float16_headdim32_causal1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-1D backward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_1D_float16_headdim64(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_1D_float16_headdim64_causal0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_1D_float16_headdim64_causal1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-1D backward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_1D_float16_headdim128(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_1D_float16_headdim128_causal0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_1D_float16_headdim128_causal1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-1D backward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_1D_bfloat16_headdim32(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_1D_bfloat16_headdim32_causal0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_1D_bfloat16_headdim32_causal1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-1D backward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_1D_bfloat16_headdim64(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_1D_bfloat16_headdim64_causal0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_1D_bfloat16_headdim64_causal1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-1D backward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_1D_bfloat16_headdim128(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_1D_bfloat16_headdim128_causal0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_1D_bfloat16_headdim128_causal1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-1D backward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim32(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim32_causal0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim32_causal0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim32_causal1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim32_causal1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim64(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim64_causal0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim64_causal0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim64_causal1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim64_causal1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim128(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim128_causal0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim128_causal0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim128_causal1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim128_causal1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim32(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim32_causal0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim32_causal0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim32_causal1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim32_causal1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim64(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim64_causal0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim64_causal0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim64_causal1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim64_causal1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim128(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim128_causal0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim128_causal0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim128_causal1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim128_causal1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim32(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim32_causal0x0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim32_causal0x0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim32_causal0x1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim32_causal0x1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim32_causal1x0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim32_causal1x0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim32_causal1x1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim32_causal1x1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim64(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim64_causal0x0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim64_causal0x0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim64_causal0x1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim64_causal0x1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim64_causal1x0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim64_causal1x0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim64_causal1x1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim64_causal1x1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim128(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim128_causal0x0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim128_causal0x0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim128_causal0x1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim128_causal0x1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim128_causal1x0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim128_causal1x0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim128_causal1x1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim128_causal1x1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim32(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim32_causal0x0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim32_causal0x0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim32_causal0x1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim32_causal0x1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim32_causal1x0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim32_causal1x0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim32_causal1x1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim32_causal1x1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim64(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim64_causal0x0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim64_causal0x0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim64_causal0x1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim64_causal0x1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim64_causal1x0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim64_causal1x0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim64_causal1x1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim64_causal1x1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim128(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim128_causal0x0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim128_causal0x0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim128_causal0x1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim128_causal0x1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim128_causal1x0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim128_causal1x0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim128_causal1x1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim128_causal1x1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fna_blackwell +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_dtype.h b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_dtype.h new file mode 100644 index 00000000..5ca7f4e1 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_dtype.h @@ -0,0 +1,68 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { +#define DISPATCH_BLACKWELL_FNA_BACKWARD_1D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (dtype == torch::kFloat16) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_1D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dtype == torch::kBFloat16) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_1D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-1D backward kernel dispatch failed! It does not support dtype " + std::string(c10::toString(dtype)) + "."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (dtype == torch::kFloat16) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dtype == torch::kBFloat16) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! It does not support dtype " + std::string(c10::toString(dtype)) + "."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (dtype == torch::kFloat16) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dtype == torch::kBFloat16) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It does not support dtype " + std::string(c10::toString(dtype)) + "."); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fna_blackwell +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_head_dim.h b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_head_dim.h new file mode 100644 index 00000000..a20ede9f --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_head_dim.h @@ -0,0 +1,125 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { +#define DISPATCH_BLACKWELL_FNA_BACKWARD_1D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (dim <= 32) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_1D_float16_headdim32(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dim <= 64) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_1D_float16_headdim64(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dim <= 128) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_1D_float16_headdim128(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA backward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for float16."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_1D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (dim <= 32) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_1D_bfloat16_headdim32(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dim <= 64) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_1D_bfloat16_headdim64(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dim <= 128) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_1D_bfloat16_headdim128(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA backward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for bfloat16."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (dim <= 32) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim32(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dim <= 64) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim64(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dim <= 128) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim128(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA backward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for float16."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (dim <= 32) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim32(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dim <= 64) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim64(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dim <= 128) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim128(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA backward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for bfloat16."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (dim <= 32) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim32(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dim <= 64) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim64(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dim <= 128) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim128(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA backward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for float16."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (dim <= 32) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim32(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dim <= 64) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim64(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dim <= 128) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim128(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA backward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for bfloat16."); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fna_blackwell +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_tile_shape.h b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_tile_shape.h new file mode 100644 index 00000000..e5178e38 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_tile_shape.h @@ -0,0 +1,2777 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { +#define DISPATCH_BLACKWELL_FNA_BACKWARD_1D_float16_headdim32_causal0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_backward_float16_128x128x32_Q128_KV128_causal0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_1D_float16_headdim32_causal1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_backward_float16_128x128x32_Q128_KV128_causal1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_1D_float16_headdim64_causal0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_backward_float16_128x128x64_Q128_KV128_causal0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_1D_float16_headdim64_causal1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_backward_float16_128x128x64_Q128_KV128_causal1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_1D_float16_headdim128_causal0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_backward_float16_128x128x128_Q128_KV128_causal0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_1D_float16_headdim128_causal1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_backward_float16_128x128x128_Q128_KV128_causal1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_1D_bfloat16_headdim32_causal0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_backward_bfloat16_128x128x32_Q128_KV128_causal0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_1D_bfloat16_headdim32_causal1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_backward_bfloat16_128x128x32_Q128_KV128_causal1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_1D_bfloat16_headdim64_causal0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_backward_bfloat16_128x128x64_Q128_KV128_causal0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_1D_bfloat16_headdim64_causal1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_backward_bfloat16_128x128x64_Q128_KV128_causal1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_1D_bfloat16_headdim128_causal0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_backward_bfloat16_128x128x128_Q128_KV128_causal0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_1D_bfloat16_headdim128_causal1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_blackwell::blackwell_fna1d_backward_bfloat16_128x128x128_Q128_KV128_causal1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-1D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim32_causal0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x32_Q16x8_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x32_Q16x8_KV8x16_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x32_Q8x16_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x32_Q8x16_KV8x16_causal0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim32_causal0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x32_Q16x8_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x32_Q16x8_KV8x16_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x32_Q8x16_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x32_Q8x16_KV8x16_causal0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim32_causal1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x32_Q16x8_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x32_Q16x8_KV8x16_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x32_Q8x16_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x32_Q8x16_KV8x16_causal1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim32_causal1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x32_Q16x8_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x32_Q16x8_KV8x16_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x32_Q8x16_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x32_Q8x16_KV8x16_causal1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim64_causal0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x64_Q16x8_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x64_Q16x8_KV8x16_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x64_Q8x16_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x64_Q8x16_KV8x16_causal0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim64_causal0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x64_Q16x8_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x64_Q16x8_KV8x16_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x64_Q8x16_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x64_Q8x16_KV8x16_causal0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim64_causal1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x64_Q16x8_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x64_Q16x8_KV8x16_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x64_Q8x16_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x64_Q8x16_KV8x16_causal1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim64_causal1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x64_Q16x8_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x64_Q16x8_KV8x16_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x64_Q8x16_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x64_Q8x16_KV8x16_causal1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim128_causal0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x128_Q16x8_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x128_Q16x8_KV8x16_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x128_Q8x16_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x128_Q8x16_KV8x16_causal0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim128_causal0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x128_Q16x8_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x128_Q16x8_KV8x16_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x128_Q8x16_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x128_Q8x16_KV8x16_causal0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim128_causal1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x128_Q16x8_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x128_Q16x8_KV8x16_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x128_Q8x16_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x128_Q8x16_KV8x16_causal1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16_headdim128_causal1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x128_Q16x8_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x128_Q16x8_KV8x16_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x128_Q8x16_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_float16_128x128x128_Q8x16_KV8x16_causal1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim32_causal0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x32_Q16x8_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x32_Q16x8_KV8x16_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x32_Q8x16_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x32_Q8x16_KV8x16_causal0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim32_causal0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x32_Q16x8_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x32_Q16x8_KV8x16_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x32_Q8x16_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x32_Q8x16_KV8x16_causal0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim32_causal1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x32_Q16x8_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x32_Q16x8_KV8x16_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x32_Q8x16_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x32_Q8x16_KV8x16_causal1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim32_causal1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x32_Q16x8_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x32_Q16x8_KV8x16_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x32_Q8x16_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x32_Q8x16_KV8x16_causal1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim64_causal0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x64_Q16x8_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x64_Q16x8_KV8x16_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x64_Q8x16_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x64_Q8x16_KV8x16_causal0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim64_causal0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x64_Q16x8_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x64_Q16x8_KV8x16_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x64_Q8x16_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x64_Q8x16_KV8x16_causal0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim64_causal1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x64_Q16x8_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x64_Q16x8_KV8x16_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x64_Q8x16_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x64_Q8x16_KV8x16_causal1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim64_causal1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x64_Q16x8_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x64_Q16x8_KV8x16_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x64_Q8x16_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x64_Q8x16_KV8x16_causal1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim128_causal0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x128_Q16x8_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x128_Q16x8_KV8x16_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x128_Q8x16_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x128_Q8x16_KV8x16_causal0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim128_causal0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x128_Q16x8_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x128_Q16x8_KV8x16_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x128_Q8x16_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x128_Q8x16_KV8x16_causal0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim128_causal1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x128_Q16x8_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x128_Q16x8_KV8x16_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x128_Q8x16_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x128_Q8x16_KV8x16_causal1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16_headdim128_causal1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x128_Q16x8_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x128_Q16x8_KV8x16_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x128_Q8x16_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna2d_backward_bfloat16_128x128x128_Q8x16_KV8x16_causal1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim32_causal0x0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q2x4x16_KV2x4x16_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x2x16_KV2x4x16_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x4x16_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim32_causal0x0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q2x4x16_KV2x4x16_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x2x16_KV2x4x16_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x4x16_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim32_causal0x1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q2x4x16_KV2x4x16_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x2x16_KV2x4x16_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x4x16_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim32_causal0x1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q2x4x16_KV2x4x16_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x2x16_KV2x4x16_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x4x16_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim32_causal1x0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q2x4x16_KV2x4x16_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x2x16_KV2x4x16_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x4x16_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim32_causal1x0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q2x4x16_KV2x4x16_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x2x16_KV2x4x16_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x4x16_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim32_causal1x1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q2x4x16_KV2x4x16_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x2x16_KV2x4x16_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x4x16_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim32_causal1x1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q2x4x16_KV2x4x16_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x2x16_KV2x4x16_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x4x16_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim64_causal0x0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q2x4x16_KV2x4x16_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x2x16_KV2x4x16_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x4x16_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim64_causal0x0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q2x4x16_KV2x4x16_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x2x16_KV2x4x16_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x4x16_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim64_causal0x1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q2x4x16_KV2x4x16_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x2x16_KV2x4x16_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x4x16_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim64_causal0x1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q2x4x16_KV2x4x16_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x2x16_KV2x4x16_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x4x16_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim64_causal1x0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q2x4x16_KV2x4x16_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x2x16_KV2x4x16_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x4x16_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim64_causal1x0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q2x4x16_KV2x4x16_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x2x16_KV2x4x16_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x4x16_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim64_causal1x1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q2x4x16_KV2x4x16_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x2x16_KV2x4x16_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x4x16_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim64_causal1x1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q2x4x16_KV2x4x16_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x2x16_KV2x4x16_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x4x16_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim128_causal0x0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q2x4x16_KV2x4x16_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x2x16_KV2x4x16_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x4x16_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim128_causal0x0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q2x4x16_KV2x4x16_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x2x16_KV2x4x16_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x4x16_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim128_causal0x1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q2x4x16_KV2x4x16_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x2x16_KV2x4x16_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x4x16_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim128_causal0x1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q2x4x16_KV2x4x16_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x2x16_KV2x4x16_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x4x16_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim128_causal1x0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q2x4x16_KV2x4x16_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x2x16_KV2x4x16_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x4x16_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim128_causal1x0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q2x4x16_KV2x4x16_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x2x16_KV2x4x16_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x4x16_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim128_causal1x1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q2x4x16_KV2x4x16_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x2x16_KV2x4x16_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x4x16_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16_headdim128_causal1x1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q2x4x16_KV2x4x16_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x2x16_KV2x4x16_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x4x16_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim32_causal0x0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q2x4x16_KV2x4x16_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x2x16_KV2x4x16_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x4x16_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim32_causal0x0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q2x4x16_KV2x4x16_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x2x16_KV2x4x16_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x4x16_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim32_causal0x1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q2x4x16_KV2x4x16_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x2x16_KV2x4x16_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x4x16_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim32_causal0x1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q2x4x16_KV2x4x16_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x2x16_KV2x4x16_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x4x16_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim32_causal1x0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q2x4x16_KV2x4x16_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x2x16_KV2x4x16_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x4x16_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim32_causal1x0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q2x4x16_KV2x4x16_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x2x16_KV2x4x16_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x4x16_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim32_causal1x1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q2x4x16_KV2x4x16_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x2x16_KV2x4x16_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x4x16_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim32_causal1x1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q2x4x16_KV2x4x16_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x2x16_KV2x4x16_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x4x16_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim64_causal0x0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q2x4x16_KV2x4x16_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x2x16_KV2x4x16_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x4x16_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim64_causal0x0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q2x4x16_KV2x4x16_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x2x16_KV2x4x16_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x4x16_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim64_causal0x1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q2x4x16_KV2x4x16_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x2x16_KV2x4x16_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x4x16_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim64_causal0x1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q2x4x16_KV2x4x16_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x2x16_KV2x4x16_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x4x16_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim64_causal1x0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q2x4x16_KV2x4x16_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x2x16_KV2x4x16_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x4x16_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim64_causal1x0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q2x4x16_KV2x4x16_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x2x16_KV2x4x16_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x4x16_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim64_causal1x1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q2x4x16_KV2x4x16_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x2x16_KV2x4x16_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x4x16_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim64_causal1x1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q2x4x16_KV2x4x16_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x2x16_KV2x4x16_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x4x16_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim128_causal0x0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q2x4x16_KV2x4x16_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x2x16_KV2x4x16_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x4x16_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim128_causal0x0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q2x4x16_KV2x4x16_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x2x16_KV2x4x16_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x4x16_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim128_causal0x1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q2x4x16_KV2x4x16_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x2x16_KV2x4x16_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x4x16_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim128_causal0x1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q2x4x16_KV2x4x16_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x2x16_KV2x4x16_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x4x16_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim128_causal1x0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q2x4x16_KV2x4x16_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x2x16_KV2x4x16_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x4x16_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim128_causal1x0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q2x4x16_KV2x4x16_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x2x16_KV2x4x16_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x4x16_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim128_causal1x1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q2x4x16_KV2x4x16_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x2x16_KV2x4x16_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x4x16_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16_headdim128_causal1x1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q2x4x16_KV2x4x16_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 2 && cute::get<2>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x2x16_KV2x4x16_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 16) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x4x16_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_blackwell::blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Blackwell FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fna_blackwell +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/interface.h b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/interface.h new file mode 100644 index 00000000..b7b59f47 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/interface.h @@ -0,0 +1,45 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { +#define DISPATCH_BLACKWELL_FNA_BACKWARD(rank, dtype, dim, is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if constexpr (rank == 1) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_1D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if constexpr (rank == 2) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_2D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if constexpr (rank == 3) { \ + DISPATCH_BLACKWELL_FNA_BACKWARD_3D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Blackwell FNA backward kernel dispatch failed! It only supports NA1D, 2D, and 3D!"); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fna_blackwell +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/kernels.h b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/kernels.h new file mode 100644 index 00000000..4ea8f8f5 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/kernels.h @@ -0,0 +1,15688 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + +void blackwell_fna1d_backward_float16_128x128x32_Q128_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_backward_float16_128x128x32_Q128_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_backward_float16_128x128x64_Q128_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_backward_float16_128x128x64_Q128_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_backward_float16_128x128x128_Q128_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_backward_float16_128x128x128_Q128_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_backward_bfloat16_128x128x32_Q128_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_backward_bfloat16_128x128x32_Q128_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_backward_bfloat16_128x128x64_Q128_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_backward_bfloat16_128x128x64_Q128_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_backward_bfloat16_128x128x128_Q128_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna1d_backward_bfloat16_128x128x128_Q128_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x32_Q16x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x32_Q16x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x32_Q8x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x32_Q8x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x32_Q16x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x32_Q16x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x32_Q8x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x32_Q8x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x32_Q16x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x32_Q16x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x32_Q8x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x32_Q8x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x32_Q16x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x32_Q16x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x32_Q8x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x32_Q8x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x64_Q16x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x64_Q16x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x64_Q8x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x64_Q8x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x64_Q16x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x64_Q16x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x64_Q8x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x64_Q8x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x64_Q16x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x64_Q16x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x64_Q8x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x64_Q8x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x64_Q16x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x64_Q16x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x64_Q8x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x64_Q8x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x128_Q16x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x128_Q16x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x128_Q8x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x128_Q8x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x128_Q16x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x128_Q16x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x128_Q8x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x128_Q8x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x128_Q16x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x128_Q16x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x128_Q8x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x128_Q8x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x128_Q16x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x128_Q16x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x128_Q8x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_float16_128x128x128_Q8x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q16x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q16x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q8x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q8x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q16x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q16x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q8x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q8x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q16x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q16x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q8x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q8x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q16x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q16x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q8x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q8x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q16x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q16x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q8x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q8x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q16x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q16x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q8x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q8x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q16x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q16x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q8x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q8x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q16x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q16x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q8x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q8x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q16x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q16x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q8x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q8x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q16x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q16x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q8x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q8x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q16x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q16x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q8x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q8x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q16x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q16x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q8x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q8x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q2x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x2x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q2x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x2x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q2x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x2x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q2x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x2x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q2x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x2x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q2x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x2x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q2x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x2x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q2x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x2x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q2x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x2x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q2x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x2x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q2x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x2x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q2x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x2x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q2x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x2x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q2x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x2x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q2x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x2x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q2x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x2x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q2x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x2x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q2x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x2x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q2x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x2x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q2x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x2x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q2x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x2x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q2x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x2x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q2x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x2x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q2x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x2x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x2x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x2x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x2x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x2x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x2x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x2x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x2x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x2x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x2x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x2x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x2x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x2x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x2x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x2x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x2x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x2x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x2x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x2x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x2x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x2x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x2x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x2x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x2x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x2x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + + +} // namespace natten +} // namespace cuda +} // namespace fna_blackwell +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/fmha/dispatch_dtype.h b/natten/csrc/autogen/include/natten_autogen/cuda/fmha/dispatch_dtype.h new file mode 100644 index 00000000..c11a5dc1 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/fmha/dispatch_dtype.h @@ -0,0 +1,129 @@ +#pragma once + + +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha { +#define DISPATCH_FMHA_FORWARD_SM50(dtype, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + fmha_sm50_float32(cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + fmha_sm50_float16(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FMHA does not support this data type on SM50."); \ + } \ +}(); + +#define DISPATCH_FMHA_FORWARD_SM70(dtype, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + fmha_sm70_float32(cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + fmha_sm70_float16(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FMHA does not support this data type on SM70."); \ + } \ +}(); + +#define DISPATCH_FMHA_FORWARD_SM75(dtype, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + fmha_sm75_float32(cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + fmha_sm75_float16(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FMHA does not support this data type on SM75."); \ + } \ +}(); + +#define DISPATCH_FMHA_FORWARD_SM80(dtype, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + fmha_sm80_float32(cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + fmha_sm80_float16(cb); \ + } \ + else if (dtype == torch::kBFloat16) { \ + fmha_sm80_bfloat16(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FMHA does not support this data type on SM80."); \ + } \ +}(); + +#define DISPATCH_FMHA_BACKWARD_SM50(dtype, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + fmha_backward_sm50_float32(cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + fmha_backward_sm50_float16(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FMHA does not support this data type on SM50."); \ + } \ +}(); + +#define DISPATCH_FMHA_BACKWARD_SM70(dtype, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + fmha_backward_sm70_float32(cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + fmha_backward_sm70_float16(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FMHA does not support this data type on SM70."); \ + } \ +}(); + +#define DISPATCH_FMHA_BACKWARD_SM75(dtype, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + fmha_backward_sm75_float32(cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + fmha_backward_sm75_float16(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FMHA does not support this data type on SM75."); \ + } \ +}(); + +#define DISPATCH_FMHA_BACKWARD_SM80(dtype, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + fmha_backward_sm80_float32(cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + fmha_backward_sm80_float16(cb); \ + } \ + else if (dtype == torch::kBFloat16) { \ + fmha_backward_sm80_bfloat16(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FMHA does not support this data type on SM80."); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fmha + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/fmha/interface.h b/natten/csrc/autogen/include/natten_autogen/cuda/fmha/interface.h new file mode 100644 index 00000000..b14ee842 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/fmha/interface.h @@ -0,0 +1,57 @@ +#pragma once + + +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha { +#define DISPATCH_FMHA_FORWARD(cc, dtype, cb) \ + [&] { \ + if (cc < 70 && cc >= 50) { \ + DISPATCH_FMHA_FORWARD_SM50(dtype, cb); \ + } \ + else if (cc < 75 && cc >= 70) { \ + DISPATCH_FMHA_FORWARD_SM70(dtype, cb); \ + } \ + else if (cc < 80 && cc >= 75) { \ + DISPATCH_FMHA_FORWARD_SM75(dtype, cb); \ + } \ + else if (cc >= 80) { \ + DISPATCH_FMHA_FORWARD_SM80(dtype, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FMHA kernel dispatch failed! FMHA is not implemented for this device."); \ + } \ +}(); + +#define DISPATCH_FMHA_BACKWARD(cc, dtype, cb) \ + [&] { \ + if (cc < 70 && cc >= 50) { \ + DISPATCH_FMHA_BACKWARD_SM50(dtype, cb); \ + } \ + else if (cc < 75 && cc >= 70) { \ + DISPATCH_FMHA_BACKWARD_SM70(dtype, cb); \ + } \ + else if (cc < 80 && cc >= 75) { \ + DISPATCH_FMHA_BACKWARD_SM75(dtype, cb); \ + } \ + else if (cc >= 80) { \ + DISPATCH_FMHA_BACKWARD_SM80(dtype, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FMHA kernel dispatch failed! FMHA is not implemented for this device."); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fmha + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/fmha/kernels.h b/natten/csrc/autogen/include/natten_autogen/cuda/fmha/kernels.h new file mode 100644 index 00000000..1d417af0 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/fmha/kernels.h @@ -0,0 +1,1226 @@ +#pragma once + + +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha { + + +/////////////////////////////////////////////////////////////////// +// FMHA / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x32_sm50_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x32_sm50_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x32_sm50_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x64_sm50_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x64_sm50_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x64_sm50_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x128_sm50_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x128_sm50_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x128_sm50_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x65536_sm50_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x65536_sm50_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x65536_sm50_float32(typename AttentionKernel::Params p); + + +template +void fmha_sm50_float32(T cb) { + cb(AttentionKernel(), fmha_32x128x32_sm50_float32); + cb(AttentionKernel(), fmha_64x64x32_sm50_float32); + cb(AttentionKernel(), fmha_64x128x32_sm50_float32); + cb(AttentionKernel(), fmha_32x128x64_sm50_float32); + cb(AttentionKernel(), fmha_64x64x64_sm50_float32); + cb(AttentionKernel(), fmha_64x128x64_sm50_float32); + cb(AttentionKernel(), fmha_32x128x128_sm50_float32); + cb(AttentionKernel(), fmha_64x64x128_sm50_float32); + cb(AttentionKernel(), fmha_64x128x128_sm50_float32); + cb(AttentionKernel(), fmha_32x128x65536_sm50_float32); + cb(AttentionKernel(), fmha_64x128x65536_sm50_float32); + cb(AttentionKernel(), fmha_64x64x65536_sm50_float32); +} + +/////////////////////////////////////////////////////////////////// +// FMHA / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x32_sm50_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x32_sm50_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x32_sm50_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x64_sm50_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x64_sm50_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x64_sm50_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x128_sm50_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x128_sm50_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x128_sm50_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x65536_sm50_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x65536_sm50_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x65536_sm50_float16(typename AttentionKernel::Params p); + + +template +void fmha_sm50_float16(T cb) { + cb(AttentionKernel(), fmha_32x128x32_sm50_float16); + cb(AttentionKernel(), fmha_64x64x32_sm50_float16); + cb(AttentionKernel(), fmha_64x128x32_sm50_float16); + cb(AttentionKernel(), fmha_32x128x64_sm50_float16); + cb(AttentionKernel(), fmha_64x64x64_sm50_float16); + cb(AttentionKernel(), fmha_64x128x64_sm50_float16); + cb(AttentionKernel(), fmha_32x128x128_sm50_float16); + cb(AttentionKernel(), fmha_64x64x128_sm50_float16); + cb(AttentionKernel(), fmha_64x128x128_sm50_float16); + cb(AttentionKernel(), fmha_32x128x65536_sm50_float16); + cb(AttentionKernel(), fmha_64x128x65536_sm50_float16); + cb(AttentionKernel(), fmha_64x64x65536_sm50_float16); +} + +/////////////////////////////////////////////////////////////////// +// FMHA / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x32_sm70_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x32_sm70_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x32_sm70_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x64_sm70_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x64_sm70_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x64_sm70_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x128_sm70_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x128_sm70_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x128_sm70_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x65536_sm70_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x65536_sm70_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x65536_sm70_float32(typename AttentionKernel::Params p); + + +template +void fmha_sm70_float32(T cb) { + cb(AttentionKernel(), fmha_32x128x32_sm70_float32); + cb(AttentionKernel(), fmha_64x64x32_sm70_float32); + cb(AttentionKernel(), fmha_64x128x32_sm70_float32); + cb(AttentionKernel(), fmha_32x128x64_sm70_float32); + cb(AttentionKernel(), fmha_64x64x64_sm70_float32); + cb(AttentionKernel(), fmha_64x128x64_sm70_float32); + cb(AttentionKernel(), fmha_32x128x128_sm70_float32); + cb(AttentionKernel(), fmha_64x64x128_sm70_float32); + cb(AttentionKernel(), fmha_64x128x128_sm70_float32); + cb(AttentionKernel(), fmha_32x128x65536_sm70_float32); + cb(AttentionKernel(), fmha_64x128x65536_sm70_float32); + cb(AttentionKernel(), fmha_64x64x65536_sm70_float32); +} + +/////////////////////////////////////////////////////////////////// +// FMHA / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x32_sm70_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x32_sm70_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x32_sm70_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x64_sm70_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x64_sm70_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x64_sm70_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x128_sm70_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x128_sm70_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x128_sm70_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x65536_sm70_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x65536_sm70_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x65536_sm70_float16(typename AttentionKernel::Params p); + + +template +void fmha_sm70_float16(T cb) { + cb(AttentionKernel(), fmha_32x128x32_sm70_float16); + cb(AttentionKernel(), fmha_64x64x32_sm70_float16); + cb(AttentionKernel(), fmha_64x128x32_sm70_float16); + cb(AttentionKernel(), fmha_32x128x64_sm70_float16); + cb(AttentionKernel(), fmha_64x64x64_sm70_float16); + cb(AttentionKernel(), fmha_64x128x64_sm70_float16); + cb(AttentionKernel(), fmha_32x128x128_sm70_float16); + cb(AttentionKernel(), fmha_64x64x128_sm70_float16); + cb(AttentionKernel(), fmha_64x128x128_sm70_float16); + cb(AttentionKernel(), fmha_32x128x65536_sm70_float16); + cb(AttentionKernel(), fmha_64x128x65536_sm70_float16); + cb(AttentionKernel(), fmha_64x64x65536_sm70_float16); +} + +/////////////////////////////////////////////////////////////////// +// FMHA / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x32_sm75_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x32_sm75_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x32_sm75_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x64_sm75_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x64_sm75_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x64_sm75_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x128_sm75_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x128_sm75_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x128_sm75_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x65536_sm75_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x65536_sm75_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x65536_sm75_float32(typename AttentionKernel::Params p); + + +template +void fmha_sm75_float32(T cb) { + cb(AttentionKernel(), fmha_32x128x32_sm75_float32); + cb(AttentionKernel(), fmha_64x64x32_sm75_float32); + cb(AttentionKernel(), fmha_64x128x32_sm75_float32); + cb(AttentionKernel(), fmha_32x128x64_sm75_float32); + cb(AttentionKernel(), fmha_64x64x64_sm75_float32); + cb(AttentionKernel(), fmha_64x128x64_sm75_float32); + cb(AttentionKernel(), fmha_32x128x128_sm75_float32); + cb(AttentionKernel(), fmha_64x64x128_sm75_float32); + cb(AttentionKernel(), fmha_64x128x128_sm75_float32); + cb(AttentionKernel(), fmha_32x128x65536_sm75_float32); + cb(AttentionKernel(), fmha_64x128x65536_sm75_float32); + cb(AttentionKernel(), fmha_64x64x65536_sm75_float32); +} + +/////////////////////////////////////////////////////////////////// +// FMHA / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x32_sm75_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x32_sm75_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x32_sm75_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x64_sm75_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x64_sm75_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x64_sm75_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x128_sm75_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x128_sm75_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x128_sm75_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x65536_sm75_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x65536_sm75_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x65536_sm75_float16(typename AttentionKernel::Params p); + + +template +void fmha_sm75_float16(T cb) { + cb(AttentionKernel(), fmha_32x128x32_sm75_float16); + cb(AttentionKernel(), fmha_64x64x32_sm75_float16); + cb(AttentionKernel(), fmha_64x128x32_sm75_float16); + cb(AttentionKernel(), fmha_32x128x64_sm75_float16); + cb(AttentionKernel(), fmha_64x64x64_sm75_float16); + cb(AttentionKernel(), fmha_64x128x64_sm75_float16); + cb(AttentionKernel(), fmha_32x128x128_sm75_float16); + cb(AttentionKernel(), fmha_64x64x128_sm75_float16); + cb(AttentionKernel(), fmha_64x128x128_sm75_float16); + cb(AttentionKernel(), fmha_32x128x65536_sm75_float16); + cb(AttentionKernel(), fmha_64x128x65536_sm75_float16); + cb(AttentionKernel(), fmha_64x64x65536_sm75_float16); +} + +/////////////////////////////////////////////////////////////////// +// FMHA / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x32_sm80_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x32_sm80_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x32_sm80_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x64_sm80_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x64_sm80_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x64_sm80_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x128_sm80_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x128_sm80_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x128_sm80_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x65536_sm80_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x65536_sm80_float32(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x65536_sm80_float32(typename AttentionKernel::Params p); + + +template +void fmha_sm80_float32(T cb) { + cb(AttentionKernel(), fmha_32x128x32_sm80_float32); + cb(AttentionKernel(), fmha_64x64x32_sm80_float32); + cb(AttentionKernel(), fmha_64x128x32_sm80_float32); + cb(AttentionKernel(), fmha_32x128x64_sm80_float32); + cb(AttentionKernel(), fmha_64x64x64_sm80_float32); + cb(AttentionKernel(), fmha_64x128x64_sm80_float32); + cb(AttentionKernel(), fmha_32x128x128_sm80_float32); + cb(AttentionKernel(), fmha_64x64x128_sm80_float32); + cb(AttentionKernel(), fmha_64x128x128_sm80_float32); + cb(AttentionKernel(), fmha_32x128x65536_sm80_float32); + cb(AttentionKernel(), fmha_64x128x65536_sm80_float32); + cb(AttentionKernel(), fmha_64x64x65536_sm80_float32); +} + +/////////////////////////////////////////////////////////////////// +// FMHA / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x32_sm80_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x32_sm80_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x32_sm80_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x64_sm80_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x64_sm80_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x64_sm80_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x128_sm80_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x128_sm80_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x128_sm80_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x65536_sm80_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x65536_sm80_float16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x65536_sm80_float16(typename AttentionKernel::Params p); + + +template +void fmha_sm80_float16(T cb) { + cb(AttentionKernel(), fmha_32x128x32_sm80_float16); + cb(AttentionKernel(), fmha_64x64x32_sm80_float16); + cb(AttentionKernel(), fmha_64x128x32_sm80_float16); + cb(AttentionKernel(), fmha_32x128x64_sm80_float16); + cb(AttentionKernel(), fmha_64x64x64_sm80_float16); + cb(AttentionKernel(), fmha_64x128x64_sm80_float16); + cb(AttentionKernel(), fmha_32x128x128_sm80_float16); + cb(AttentionKernel(), fmha_64x64x128_sm80_float16); + cb(AttentionKernel(), fmha_64x128x128_sm80_float16); + cb(AttentionKernel(), fmha_32x128x65536_sm80_float16); + cb(AttentionKernel(), fmha_64x128x65536_sm80_float16); + cb(AttentionKernel(), fmha_64x64x65536_sm80_float16); +} + +/////////////////////////////////////////////////////////////////// +// FMHA / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x32_sm80_bfloat16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x32_sm80_bfloat16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x32_sm80_bfloat16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x64_sm80_bfloat16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x64_sm80_bfloat16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x64_sm80_bfloat16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x128_sm80_bfloat16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x128_sm80_bfloat16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x128_sm80_bfloat16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x65536_sm80_bfloat16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x65536_sm80_bfloat16(typename AttentionKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x65536_sm80_bfloat16(typename AttentionKernel::Params p); + + +template +void fmha_sm80_bfloat16(T cb) { + cb(AttentionKernel(), fmha_32x128x32_sm80_bfloat16); + cb(AttentionKernel(), fmha_64x64x32_sm80_bfloat16); + cb(AttentionKernel(), fmha_64x128x32_sm80_bfloat16); + cb(AttentionKernel(), fmha_32x128x64_sm80_bfloat16); + cb(AttentionKernel(), fmha_64x64x64_sm80_bfloat16); + cb(AttentionKernel(), fmha_64x128x64_sm80_bfloat16); + cb(AttentionKernel(), fmha_32x128x128_sm80_bfloat16); + cb(AttentionKernel(), fmha_64x64x128_sm80_bfloat16); + cb(AttentionKernel(), fmha_64x128x128_sm80_bfloat16); + cb(AttentionKernel(), fmha_32x128x65536_sm80_bfloat16); + cb(AttentionKernel(), fmha_64x128x65536_sm80_bfloat16); + cb(AttentionKernel(), fmha_64x64x65536_sm80_bfloat16); +} + +/////////////////////////////////////////////////////////////////// +// FMHA / float32 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x32_sm50_float32(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x64_sm50_float32(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x128_sm50_float32(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x65536_sm50_float32(typename AttentionBackwardKernel::Params p); + + +template +void fmha_backward_sm50_float32(T cb) { + cb(AttentionBackwardKernel(), fmha_backward_64x64x32_sm50_float32); + cb(AttentionBackwardKernel(), fmha_backward_64x64x64_sm50_float32); + cb(AttentionBackwardKernel(), fmha_backward_64x64x128_sm50_float32); + cb(AttentionBackwardKernel(), fmha_backward_64x64x65536_sm50_float32); +} + +/////////////////////////////////////////////////////////////////// +// FMHA / float16 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x32_sm50_float16(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x64_sm50_float16(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x128_sm50_float16(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x65536_sm50_float16(typename AttentionBackwardKernel::Params p); + + +template +void fmha_backward_sm50_float16(T cb) { + cb(AttentionBackwardKernel(), fmha_backward_64x64x32_sm50_float16); + cb(AttentionBackwardKernel(), fmha_backward_64x64x64_sm50_float16); + cb(AttentionBackwardKernel(), fmha_backward_64x64x128_sm50_float16); + cb(AttentionBackwardKernel(), fmha_backward_64x64x65536_sm50_float16); +} + +/////////////////////////////////////////////////////////////////// +// FMHA / float32 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x32_sm70_float32(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x64_sm70_float32(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x128_sm70_float32(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x65536_sm70_float32(typename AttentionBackwardKernel::Params p); + + +template +void fmha_backward_sm70_float32(T cb) { + cb(AttentionBackwardKernel(), fmha_backward_64x64x32_sm70_float32); + cb(AttentionBackwardKernel(), fmha_backward_64x64x64_sm70_float32); + cb(AttentionBackwardKernel(), fmha_backward_64x64x128_sm70_float32); + cb(AttentionBackwardKernel(), fmha_backward_64x64x65536_sm70_float32); +} + +/////////////////////////////////////////////////////////////////// +// FMHA / float16 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x32_sm70_float16(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x64_sm70_float16(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x128_sm70_float16(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_128x64x128_sm70_float16(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x65536_sm70_float16(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_128x64x65536_sm70_float16(typename AttentionBackwardKernel::Params p); + + +template +void fmha_backward_sm70_float16(T cb) { + cb(AttentionBackwardKernel(), fmha_backward_64x64x32_sm70_float16); + cb(AttentionBackwardKernel(), fmha_backward_64x64x64_sm70_float16); + cb(AttentionBackwardKernel(), fmha_backward_64x64x128_sm70_float16); + cb(AttentionBackwardKernel(), fmha_backward_128x64x128_sm70_float16); + cb(AttentionBackwardKernel(), fmha_backward_64x64x65536_sm70_float16); + cb(AttentionBackwardKernel(), fmha_backward_128x64x65536_sm70_float16); +} + +/////////////////////////////////////////////////////////////////// +// FMHA / float32 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x32_sm75_float32(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x64_sm75_float32(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x128_sm75_float32(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x65536_sm75_float32(typename AttentionBackwardKernel::Params p); + + +template +void fmha_backward_sm75_float32(T cb) { + cb(AttentionBackwardKernel(), fmha_backward_64x64x32_sm75_float32); + cb(AttentionBackwardKernel(), fmha_backward_64x64x64_sm75_float32); + cb(AttentionBackwardKernel(), fmha_backward_64x64x128_sm75_float32); + cb(AttentionBackwardKernel(), fmha_backward_64x64x65536_sm75_float32); +} + +/////////////////////////////////////////////////////////////////// +// FMHA / float16 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x32_sm75_float16(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x64_sm75_float16(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x128_sm75_float16(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x65536_sm75_float16(typename AttentionBackwardKernel::Params p); + + +template +void fmha_backward_sm75_float16(T cb) { + cb(AttentionBackwardKernel(), fmha_backward_64x64x32_sm75_float16); + cb(AttentionBackwardKernel(), fmha_backward_64x64x64_sm75_float16); + cb(AttentionBackwardKernel(), fmha_backward_64x64x128_sm75_float16); + cb(AttentionBackwardKernel(), fmha_backward_64x64x65536_sm75_float16); +} + +/////////////////////////////////////////////////////////////////// +// FMHA / float32 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x32_sm80_float32(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x64_sm80_float32(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x128_sm80_float32(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_128x64x128_sm80_float32(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x65536_sm80_float32(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_128x64x65536_sm80_float32(typename AttentionBackwardKernel::Params p); + + +template +void fmha_backward_sm80_float32(T cb) { + cb(AttentionBackwardKernel(), fmha_backward_64x64x32_sm80_float32); + cb(AttentionBackwardKernel(), fmha_backward_64x64x64_sm80_float32); + cb(AttentionBackwardKernel(), fmha_backward_64x64x128_sm80_float32); + cb(AttentionBackwardKernel(), fmha_backward_128x64x128_sm80_float32); + cb(AttentionBackwardKernel(), fmha_backward_64x64x65536_sm80_float32); + cb(AttentionBackwardKernel(), fmha_backward_128x64x65536_sm80_float32); +} + +/////////////////////////////////////////////////////////////////// +// FMHA / float16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x32_sm80_float16(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x64_sm80_float16(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x128_sm80_float16(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_128x128x128_sm80_float16(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x65536_sm80_float16(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_128x64x65536_sm80_float16(typename AttentionBackwardKernel::Params p); + + +template +void fmha_backward_sm80_float16(T cb) { + cb(AttentionBackwardKernel(), fmha_backward_64x64x32_sm80_float16); + cb(AttentionBackwardKernel(), fmha_backward_64x64x64_sm80_float16); + cb(AttentionBackwardKernel(), fmha_backward_64x64x128_sm80_float16); + cb(AttentionBackwardKernel(), fmha_backward_128x128x128_sm80_float16); + cb(AttentionBackwardKernel(), fmha_backward_64x64x65536_sm80_float16); + cb(AttentionBackwardKernel(), fmha_backward_128x64x65536_sm80_float16); +} + +/////////////////////////////////////////////////////////////////// +// FMHA / bfloat16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x32_sm80_bfloat16(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x64_sm80_bfloat16(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x128_sm80_bfloat16(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_128x128x128_sm80_bfloat16(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x65536_sm80_bfloat16(typename AttentionBackwardKernel::Params p); + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_128x64x65536_sm80_bfloat16(typename AttentionBackwardKernel::Params p); + + +template +void fmha_backward_sm80_bfloat16(T cb) { + cb(AttentionBackwardKernel(), fmha_backward_64x64x32_sm80_bfloat16); + cb(AttentionBackwardKernel(), fmha_backward_64x64x64_sm80_bfloat16); + cb(AttentionBackwardKernel(), fmha_backward_64x64x128_sm80_bfloat16); + cb(AttentionBackwardKernel(), fmha_backward_128x128x128_sm80_bfloat16); + cb(AttentionBackwardKernel(), fmha_backward_64x64x65536_sm80_bfloat16); + cb(AttentionBackwardKernel(), fmha_backward_128x64x65536_sm80_bfloat16); +} + +} // namespace natten +} // namespace cuda +} // namespace fmha + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/fna/dispatch_cm.h b/natten/csrc/autogen/include/natten_autogen/cuda/fna/dispatch_cm.h new file mode 100644 index 00000000..ebf3ae86 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/fna/dispatch_cm.h @@ -0,0 +1,1154 @@ +#pragma once + + +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { +#define DISPATCH_FNA_FORWARD_1D_SM50_float32(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal)) { \ + fna1d_sm50_float32_cm_0(cb); \ + } \ + else if (std::get<0>(is_causal)) { \ + fna1d_sm50_float32_cm_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-1D, SM50, float32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_1D_SM50_float16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal)) { \ + fna1d_sm50_float16_cm_0(cb); \ + } \ + else if (std::get<0>(is_causal)) { \ + fna1d_sm50_float16_cm_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-1D, SM50, float16) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_1D_SM70_float32(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal)) { \ + fna1d_sm70_float32_cm_0(cb); \ + } \ + else if (std::get<0>(is_causal)) { \ + fna1d_sm70_float32_cm_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-1D, SM70, float32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_1D_SM70_float16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal)) { \ + fna1d_sm70_float16_cm_0(cb); \ + } \ + else if (std::get<0>(is_causal)) { \ + fna1d_sm70_float16_cm_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-1D, SM70, float16) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_1D_SM75_float32(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal)) { \ + fna1d_sm75_float32_cm_0(cb); \ + } \ + else if (std::get<0>(is_causal)) { \ + fna1d_sm75_float32_cm_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-1D, SM75, float32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_1D_SM75_float16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal)) { \ + fna1d_sm75_float16_cm_0(cb); \ + } \ + else if (std::get<0>(is_causal)) { \ + fna1d_sm75_float16_cm_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-1D, SM75, float16) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_1D_SM80_float32(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal)) { \ + fna1d_sm80_float32_cm_0(cb); \ + } \ + else if (std::get<0>(is_causal)) { \ + fna1d_sm80_float32_cm_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-1D, SM80, float32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_1D_SM80_float16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal)) { \ + fna1d_sm80_float16_cm_0(cb); \ + } \ + else if (std::get<0>(is_causal)) { \ + fna1d_sm80_float16_cm_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-1D, SM80, float16) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_1D_SM80_bfloat16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal)) { \ + fna1d_sm80_bfloat16_cm_0(cb); \ + } \ + else if (std::get<0>(is_causal)) { \ + fna1d_sm80_bfloat16_cm_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-1D, SM80, bfloat16) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_2D_SM50_float32(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_sm50_float32_cm_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_sm50_float32_cm_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_sm50_float32_cm_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_sm50_float32_cm_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-2D, SM50, float32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_2D_SM50_float16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_sm50_float16_cm_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_sm50_float16_cm_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_sm50_float16_cm_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_sm50_float16_cm_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-2D, SM50, float16) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_2D_SM70_float32(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_sm70_float32_cm_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_sm70_float32_cm_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_sm70_float32_cm_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_sm70_float32_cm_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-2D, SM70, float32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_2D_SM70_float16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_sm70_float16_cm_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_sm70_float16_cm_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_sm70_float16_cm_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_sm70_float16_cm_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-2D, SM70, float16) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_2D_SM75_float32(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_sm75_float32_cm_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_sm75_float32_cm_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_sm75_float32_cm_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_sm75_float32_cm_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-2D, SM75, float32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_2D_SM75_float16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_sm75_float16_cm_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_sm75_float16_cm_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_sm75_float16_cm_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_sm75_float16_cm_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-2D, SM75, float16) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_2D_SM80_float32(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_sm80_float32_cm_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_sm80_float32_cm_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_sm80_float32_cm_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_sm80_float32_cm_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-2D, SM80, float32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_2D_SM80_float16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_sm80_float16_cm_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_sm80_float16_cm_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_sm80_float16_cm_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_sm80_float16_cm_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-2D, SM80, float16) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_2D_SM80_bfloat16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_sm80_bfloat16_cm_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_sm80_bfloat16_cm_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_sm80_bfloat16_cm_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_sm80_bfloat16_cm_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-2D, SM80, bfloat16) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_3D_SM50_float32(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm50_float32_cm_0_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm50_float32_cm_0_0_1(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm50_float32_cm_0_1_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm50_float32_cm_0_1_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm50_float32_cm_1_0_0(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm50_float32_cm_1_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm50_float32_cm_1_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm50_float32_cm_1_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-3D, SM50, float32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_3D_SM50_float16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm50_float16_cm_0_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm50_float16_cm_0_0_1(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm50_float16_cm_0_1_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm50_float16_cm_0_1_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm50_float16_cm_1_0_0(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm50_float16_cm_1_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm50_float16_cm_1_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm50_float16_cm_1_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-3D, SM50, float16) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_3D_SM70_float32(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm70_float32_cm_0_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm70_float32_cm_0_0_1(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm70_float32_cm_0_1_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm70_float32_cm_0_1_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm70_float32_cm_1_0_0(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm70_float32_cm_1_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm70_float32_cm_1_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm70_float32_cm_1_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-3D, SM70, float32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_3D_SM70_float16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm70_float16_cm_0_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm70_float16_cm_0_0_1(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm70_float16_cm_0_1_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm70_float16_cm_0_1_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm70_float16_cm_1_0_0(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm70_float16_cm_1_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm70_float16_cm_1_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm70_float16_cm_1_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-3D, SM70, float16) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_3D_SM75_float32(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm75_float32_cm_0_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm75_float32_cm_0_0_1(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm75_float32_cm_0_1_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm75_float32_cm_0_1_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm75_float32_cm_1_0_0(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm75_float32_cm_1_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm75_float32_cm_1_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm75_float32_cm_1_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-3D, SM75, float32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_3D_SM75_float16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm75_float16_cm_0_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm75_float16_cm_0_0_1(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm75_float16_cm_0_1_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm75_float16_cm_0_1_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm75_float16_cm_1_0_0(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm75_float16_cm_1_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm75_float16_cm_1_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm75_float16_cm_1_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-3D, SM75, float16) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_3D_SM80_float32(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm80_float32_cm_0_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm80_float32_cm_0_0_1(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm80_float32_cm_0_1_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm80_float32_cm_0_1_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm80_float32_cm_1_0_0(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm80_float32_cm_1_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm80_float32_cm_1_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm80_float32_cm_1_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-3D, SM80, float32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_3D_SM80_float16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm80_float16_cm_0_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm80_float16_cm_0_0_1(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm80_float16_cm_0_1_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm80_float16_cm_0_1_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm80_float16_cm_1_0_0(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm80_float16_cm_1_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm80_float16_cm_1_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm80_float16_cm_1_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-3D, SM80, float16) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_3D_SM80_bfloat16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm80_bfloat16_cm_0_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm80_bfloat16_cm_0_0_1(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm80_bfloat16_cm_0_1_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm80_bfloat16_cm_0_1_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm80_bfloat16_cm_1_0_0(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm80_bfloat16_cm_1_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_sm80_bfloat16_cm_1_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_sm80_bfloat16_cm_1_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-3D, SM80, bfloat16) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_1D_SM50_float32(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal)) { \ + fna1d_backward_sm50_float32_cm_0(cb); \ + } \ + else if (std::get<0>(is_causal)) { \ + fna1d_backward_sm50_float32_cm_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-1D-backward, SM50, float32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_1D_SM50_float16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal)) { \ + fna1d_backward_sm50_float16_cm_0(cb); \ + } \ + else if (std::get<0>(is_causal)) { \ + fna1d_backward_sm50_float16_cm_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-1D-backward, SM50, float16) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_1D_SM70_float32(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal)) { \ + fna1d_backward_sm70_float32_cm_0(cb); \ + } \ + else if (std::get<0>(is_causal)) { \ + fna1d_backward_sm70_float32_cm_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-1D-backward, SM70, float32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_1D_SM70_float16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal)) { \ + fna1d_backward_sm70_float16_cm_0(cb); \ + } \ + else if (std::get<0>(is_causal)) { \ + fna1d_backward_sm70_float16_cm_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-1D-backward, SM70, float16) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_1D_SM75_float32(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal)) { \ + fna1d_backward_sm75_float32_cm_0(cb); \ + } \ + else if (std::get<0>(is_causal)) { \ + fna1d_backward_sm75_float32_cm_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-1D-backward, SM75, float32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_1D_SM75_float16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal)) { \ + fna1d_backward_sm75_float16_cm_0(cb); \ + } \ + else if (std::get<0>(is_causal)) { \ + fna1d_backward_sm75_float16_cm_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-1D-backward, SM75, float16) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_1D_SM80_float32(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal)) { \ + fna1d_backward_sm80_float32_cm_0(cb); \ + } \ + else if (std::get<0>(is_causal)) { \ + fna1d_backward_sm80_float32_cm_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-1D-backward, SM80, float32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_1D_SM80_float16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal)) { \ + fna1d_backward_sm80_float16_cm_0(cb); \ + } \ + else if (std::get<0>(is_causal)) { \ + fna1d_backward_sm80_float16_cm_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-1D-backward, SM80, float16) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_1D_SM80_bfloat16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal)) { \ + fna1d_backward_sm80_bfloat16_cm_0(cb); \ + } \ + else if (std::get<0>(is_causal)) { \ + fna1d_backward_sm80_bfloat16_cm_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-1D-backward, SM80, bfloat16) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_2D_SM50_float32(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_backward_sm50_float32_cm_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_backward_sm50_float32_cm_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_backward_sm50_float32_cm_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_backward_sm50_float32_cm_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-2D-backward, SM50, float32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_2D_SM50_float16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_backward_sm50_float16_cm_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_backward_sm50_float16_cm_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_backward_sm50_float16_cm_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_backward_sm50_float16_cm_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-2D-backward, SM50, float16) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_2D_SM70_float32(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_backward_sm70_float32_cm_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_backward_sm70_float32_cm_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_backward_sm70_float32_cm_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_backward_sm70_float32_cm_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-2D-backward, SM70, float32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_2D_SM70_float16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_backward_sm70_float16_cm_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_backward_sm70_float16_cm_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_backward_sm70_float16_cm_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_backward_sm70_float16_cm_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-2D-backward, SM70, float16) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_2D_SM75_float32(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_backward_sm75_float32_cm_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_backward_sm75_float32_cm_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_backward_sm75_float32_cm_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_backward_sm75_float32_cm_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-2D-backward, SM75, float32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_2D_SM75_float16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_backward_sm75_float16_cm_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_backward_sm75_float16_cm_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_backward_sm75_float16_cm_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_backward_sm75_float16_cm_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-2D-backward, SM75, float16) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_2D_SM80_float32(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_backward_sm80_float32_cm_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_backward_sm80_float32_cm_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_backward_sm80_float32_cm_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_backward_sm80_float32_cm_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-2D-backward, SM80, float32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_2D_SM80_float16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_backward_sm80_float16_cm_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_backward_sm80_float16_cm_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_backward_sm80_float16_cm_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_backward_sm80_float16_cm_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-2D-backward, SM80, float16) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_2D_SM80_bfloat16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_backward_sm80_bfloat16_cm_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_backward_sm80_bfloat16_cm_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal)) { \ + fna2d_backward_sm80_bfloat16_cm_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal)) { \ + fna2d_backward_sm80_bfloat16_cm_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-2D-backward, SM80, bfloat16) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_3D_SM50_float32(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm50_float32_cm_0_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm50_float32_cm_0_0_1(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm50_float32_cm_0_1_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm50_float32_cm_0_1_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm50_float32_cm_1_0_0(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm50_float32_cm_1_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm50_float32_cm_1_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm50_float32_cm_1_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-3D-backward, SM50, float32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_3D_SM50_float16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm50_float16_cm_0_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm50_float16_cm_0_0_1(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm50_float16_cm_0_1_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm50_float16_cm_0_1_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm50_float16_cm_1_0_0(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm50_float16_cm_1_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm50_float16_cm_1_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm50_float16_cm_1_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-3D-backward, SM50, float16) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_3D_SM70_float32(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm70_float32_cm_0_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm70_float32_cm_0_0_1(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm70_float32_cm_0_1_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm70_float32_cm_0_1_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm70_float32_cm_1_0_0(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm70_float32_cm_1_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm70_float32_cm_1_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm70_float32_cm_1_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-3D-backward, SM70, float32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_3D_SM70_float16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm70_float16_cm_0_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm70_float16_cm_0_0_1(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm70_float16_cm_0_1_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm70_float16_cm_0_1_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm70_float16_cm_1_0_0(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm70_float16_cm_1_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm70_float16_cm_1_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm70_float16_cm_1_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-3D-backward, SM70, float16) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_3D_SM75_float32(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm75_float32_cm_0_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm75_float32_cm_0_0_1(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm75_float32_cm_0_1_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm75_float32_cm_0_1_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm75_float32_cm_1_0_0(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm75_float32_cm_1_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm75_float32_cm_1_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm75_float32_cm_1_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-3D-backward, SM75, float32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_3D_SM75_float16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm75_float16_cm_0_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm75_float16_cm_0_0_1(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm75_float16_cm_0_1_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm75_float16_cm_0_1_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm75_float16_cm_1_0_0(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm75_float16_cm_1_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm75_float16_cm_1_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm75_float16_cm_1_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-3D-backward, SM75, float16) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_3D_SM80_float32(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm80_float32_cm_0_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm80_float32_cm_0_0_1(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm80_float32_cm_0_1_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm80_float32_cm_0_1_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm80_float32_cm_1_0_0(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm80_float32_cm_1_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm80_float32_cm_1_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm80_float32_cm_1_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-3D-backward, SM80, float32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_3D_SM80_float16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm80_float16_cm_0_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm80_float16_cm_0_0_1(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm80_float16_cm_0_1_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm80_float16_cm_0_1_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm80_float16_cm_1_0_0(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm80_float16_cm_1_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm80_float16_cm_1_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm80_float16_cm_1_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-3D-backward, SM80, float16) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_3D_SM80_bfloat16(is_causal, cb) \ + [&] { \ + if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm80_bfloat16_cm_0_0_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm80_bfloat16_cm_0_0_1(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm80_bfloat16_cm_0_1_0(cb); \ + } \ + else if (!std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm80_bfloat16_cm_0_1_1(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm80_bfloat16_cm_1_0_0(cb); \ + } \ + else if (std::get<0>(is_causal) && !std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm80_bfloat16_cm_1_0_1(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && !std::get<2>(is_causal)) { \ + fna3d_backward_sm80_bfloat16_cm_1_1_0(cb); \ + } \ + else if (std::get<0>(is_causal) && std::get<1>(is_causal) && std::get<2>(is_causal)) { \ + fna3d_backward_sm80_bfloat16_cm_1_1_1(cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Causal mask dispatcher (FNA-3D-backward, SM80, bfloat16) got invalid causal mask!"); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fna + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/fna/dispatch_device.h b/natten/csrc/autogen/include/natten_autogen/cuda/fna/dispatch_device.h new file mode 100644 index 00000000..17a53a9f --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/fna/dispatch_device.h @@ -0,0 +1,134 @@ +#pragma once + + +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { +#define DISPATCH_FNA_FORWARD_1D(cc, dtype, is_causal, cb) \ + [&] { \ + if (cc < 70 && cc >= 50) { \ + DISPATCH_FNA_FORWARD_1D_SM50(dtype, is_causal, cb); \ + } \ + else if (cc < 75 && cc >= 70) { \ + DISPATCH_FNA_FORWARD_1D_SM70(dtype, is_causal, cb); \ + } \ + else if (cc < 80 && cc >= 75) { \ + DISPATCH_FNA_FORWARD_1D_SM75(dtype, is_causal, cb); \ + } \ + else if (cc >= 80) { \ + DISPATCH_FNA_FORWARD_1D_SM80(dtype, is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Fused neighborhood attention is not implemented for this device."); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_2D(cc, dtype, is_causal, cb) \ + [&] { \ + if (cc < 70 && cc >= 50) { \ + DISPATCH_FNA_FORWARD_2D_SM50(dtype, is_causal, cb); \ + } \ + else if (cc < 75 && cc >= 70) { \ + DISPATCH_FNA_FORWARD_2D_SM70(dtype, is_causal, cb); \ + } \ + else if (cc < 80 && cc >= 75) { \ + DISPATCH_FNA_FORWARD_2D_SM75(dtype, is_causal, cb); \ + } \ + else if (cc >= 80) { \ + DISPATCH_FNA_FORWARD_2D_SM80(dtype, is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Fused neighborhood attention is not implemented for this device."); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_3D(cc, dtype, is_causal, cb) \ + [&] { \ + if (cc < 70 && cc >= 50) { \ + DISPATCH_FNA_FORWARD_3D_SM50(dtype, is_causal, cb); \ + } \ + else if (cc < 75 && cc >= 70) { \ + DISPATCH_FNA_FORWARD_3D_SM70(dtype, is_causal, cb); \ + } \ + else if (cc < 80 && cc >= 75) { \ + DISPATCH_FNA_FORWARD_3D_SM75(dtype, is_causal, cb); \ + } \ + else if (cc >= 80) { \ + DISPATCH_FNA_FORWARD_3D_SM80(dtype, is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Fused neighborhood attention is not implemented for this device."); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_1D(cc, dtype, is_causal, cb) \ + [&] { \ + if (cc < 70 && cc >= 50) { \ + DISPATCH_FNA_BACKWARD_1D_SM50(dtype, is_causal, cb); \ + } \ + else if (cc < 75 && cc >= 70) { \ + DISPATCH_FNA_BACKWARD_1D_SM70(dtype, is_causal, cb); \ + } \ + else if (cc < 80 && cc >= 75) { \ + DISPATCH_FNA_BACKWARD_1D_SM75(dtype, is_causal, cb); \ + } \ + else if (cc >= 80) { \ + DISPATCH_FNA_BACKWARD_1D_SM80(dtype, is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Fused neighborhood attention is not implemented for this device."); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_2D(cc, dtype, is_causal, cb) \ + [&] { \ + if (cc < 70 && cc >= 50) { \ + DISPATCH_FNA_BACKWARD_2D_SM50(dtype, is_causal, cb); \ + } \ + else if (cc < 75 && cc >= 70) { \ + DISPATCH_FNA_BACKWARD_2D_SM70(dtype, is_causal, cb); \ + } \ + else if (cc < 80 && cc >= 75) { \ + DISPATCH_FNA_BACKWARD_2D_SM75(dtype, is_causal, cb); \ + } \ + else if (cc >= 80) { \ + DISPATCH_FNA_BACKWARD_2D_SM80(dtype, is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Fused neighborhood attention is not implemented for this device."); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_3D(cc, dtype, is_causal, cb) \ + [&] { \ + if (cc < 70 && cc >= 50) { \ + DISPATCH_FNA_BACKWARD_3D_SM50(dtype, is_causal, cb); \ + } \ + else if (cc < 75 && cc >= 70) { \ + DISPATCH_FNA_BACKWARD_3D_SM70(dtype, is_causal, cb); \ + } \ + else if (cc < 80 && cc >= 75) { \ + DISPATCH_FNA_BACKWARD_3D_SM75(dtype, is_causal, cb); \ + } \ + else if (cc >= 80) { \ + DISPATCH_FNA_BACKWARD_3D_SM80(dtype, is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Fused neighborhood attention is not implemented for this device."); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fna + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/fna/dispatch_dtype.h b/natten/csrc/autogen/include/natten_autogen/cuda/fna/dispatch_dtype.h new file mode 100644 index 00000000..cbbb324c --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/fna/dispatch_dtype.h @@ -0,0 +1,350 @@ +#pragma once + + +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { +#define DISPATCH_FNA_FORWARD_1D_SM50(dtype, is_causal, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_FNA_FORWARD_1D_SM50_float32(is_causal, cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_FNA_FORWARD_1D_SM50_float16(is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FNA-1D does not support this data type on SM50."); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_1D_SM70(dtype, is_causal, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_FNA_FORWARD_1D_SM70_float32(is_causal, cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_FNA_FORWARD_1D_SM70_float16(is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FNA-1D does not support this data type on SM70."); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_1D_SM75(dtype, is_causal, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_FNA_FORWARD_1D_SM75_float32(is_causal, cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_FNA_FORWARD_1D_SM75_float16(is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FNA-1D does not support this data type on SM75."); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_1D_SM80(dtype, is_causal, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_FNA_FORWARD_1D_SM80_float32(is_causal, cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_FNA_FORWARD_1D_SM80_float16(is_causal, cb); \ + } \ + else if (dtype == torch::kBFloat16) { \ + DISPATCH_FNA_FORWARD_1D_SM80_bfloat16(is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FNA-1D does not support this data type on SM80."); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_2D_SM50(dtype, is_causal, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_FNA_FORWARD_2D_SM50_float32(is_causal, cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_FNA_FORWARD_2D_SM50_float16(is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FNA-2D does not support this data type on SM50."); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_2D_SM70(dtype, is_causal, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_FNA_FORWARD_2D_SM70_float32(is_causal, cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_FNA_FORWARD_2D_SM70_float16(is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FNA-2D does not support this data type on SM70."); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_2D_SM75(dtype, is_causal, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_FNA_FORWARD_2D_SM75_float32(is_causal, cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_FNA_FORWARD_2D_SM75_float16(is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FNA-2D does not support this data type on SM75."); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_2D_SM80(dtype, is_causal, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_FNA_FORWARD_2D_SM80_float32(is_causal, cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_FNA_FORWARD_2D_SM80_float16(is_causal, cb); \ + } \ + else if (dtype == torch::kBFloat16) { \ + DISPATCH_FNA_FORWARD_2D_SM80_bfloat16(is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FNA-2D does not support this data type on SM80."); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_3D_SM50(dtype, is_causal, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_FNA_FORWARD_3D_SM50_float32(is_causal, cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_FNA_FORWARD_3D_SM50_float16(is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FNA-3D does not support this data type on SM50."); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_3D_SM70(dtype, is_causal, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_FNA_FORWARD_3D_SM70_float32(is_causal, cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_FNA_FORWARD_3D_SM70_float16(is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FNA-3D does not support this data type on SM70."); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_3D_SM75(dtype, is_causal, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_FNA_FORWARD_3D_SM75_float32(is_causal, cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_FNA_FORWARD_3D_SM75_float16(is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FNA-3D does not support this data type on SM75."); \ + } \ +}(); + +#define DISPATCH_FNA_FORWARD_3D_SM80(dtype, is_causal, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_FNA_FORWARD_3D_SM80_float32(is_causal, cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_FNA_FORWARD_3D_SM80_float16(is_causal, cb); \ + } \ + else if (dtype == torch::kBFloat16) { \ + DISPATCH_FNA_FORWARD_3D_SM80_bfloat16(is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FNA-3D does not support this data type on SM80."); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_1D_SM50(dtype, is_causal, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_FNA_BACKWARD_1D_SM50_float32(is_causal, cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_FNA_BACKWARD_1D_SM50_float16(is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FNA-1D does not support this data type on SM50."); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_1D_SM70(dtype, is_causal, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_FNA_BACKWARD_1D_SM70_float32(is_causal, cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_FNA_BACKWARD_1D_SM70_float16(is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FNA-1D does not support this data type on SM70."); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_1D_SM75(dtype, is_causal, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_FNA_BACKWARD_1D_SM75_float32(is_causal, cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_FNA_BACKWARD_1D_SM75_float16(is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FNA-1D does not support this data type on SM75."); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_1D_SM80(dtype, is_causal, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_FNA_BACKWARD_1D_SM80_float32(is_causal, cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_FNA_BACKWARD_1D_SM80_float16(is_causal, cb); \ + } \ + else if (dtype == torch::kBFloat16) { \ + DISPATCH_FNA_BACKWARD_1D_SM80_bfloat16(is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FNA-1D does not support this data type on SM80."); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_2D_SM50(dtype, is_causal, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_FNA_BACKWARD_2D_SM50_float32(is_causal, cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_FNA_BACKWARD_2D_SM50_float16(is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FNA-2D does not support this data type on SM50."); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_2D_SM70(dtype, is_causal, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_FNA_BACKWARD_2D_SM70_float32(is_causal, cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_FNA_BACKWARD_2D_SM70_float16(is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FNA-2D does not support this data type on SM70."); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_2D_SM75(dtype, is_causal, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_FNA_BACKWARD_2D_SM75_float32(is_causal, cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_FNA_BACKWARD_2D_SM75_float16(is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FNA-2D does not support this data type on SM75."); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_2D_SM80(dtype, is_causal, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_FNA_BACKWARD_2D_SM80_float32(is_causal, cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_FNA_BACKWARD_2D_SM80_float16(is_causal, cb); \ + } \ + else if (dtype == torch::kBFloat16) { \ + DISPATCH_FNA_BACKWARD_2D_SM80_bfloat16(is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FNA-2D does not support this data type on SM80."); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_3D_SM50(dtype, is_causal, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_FNA_BACKWARD_3D_SM50_float32(is_causal, cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_FNA_BACKWARD_3D_SM50_float16(is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FNA-3D does not support this data type on SM50."); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_3D_SM70(dtype, is_causal, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_FNA_BACKWARD_3D_SM70_float32(is_causal, cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_FNA_BACKWARD_3D_SM70_float16(is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FNA-3D does not support this data type on SM70."); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_3D_SM75(dtype, is_causal, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_FNA_BACKWARD_3D_SM75_float32(is_causal, cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_FNA_BACKWARD_3D_SM75_float16(is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FNA-3D does not support this data type on SM75."); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_3D_SM80(dtype, is_causal, cb) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_FNA_BACKWARD_3D_SM80_float32(is_causal, cb); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_FNA_BACKWARD_3D_SM80_float16(is_causal, cb); \ + } \ + else if (dtype == torch::kBFloat16) { \ + DISPATCH_FNA_BACKWARD_3D_SM80_bfloat16(is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN kernel dispatch failed! FNA-3D does not support this data type on SM80."); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fna + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/fna/interface.h b/natten/csrc/autogen/include/natten_autogen/cuda/fna/interface.h new file mode 100644 index 00000000..cd7a7c2a --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/fna/interface.h @@ -0,0 +1,52 @@ +#pragma once + + +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { +#define DISPATCH_FNA_FORWARD_KERNEL(rank, cc, dtype, is_causal, cb) \ + [&] { \ + if constexpr (rank == 1) { \ + DISPATCH_FNA_FORWARD_1D(cc, dtype, is_causal, cb); \ + } \ + else if constexpr (rank == 2) { \ + DISPATCH_FNA_FORWARD_2D(cc, dtype, is_causal, cb); \ + } \ + else if constexpr (rank == 3) { \ + DISPATCH_FNA_FORWARD_3D(cc, dtype, is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Invalid spatial extent rank! Only 1, 2, and 3D are supported!"); \ + } \ +}(); + +#define DISPATCH_FNA_BACKWARD_KERNEL(rank, cc, dtype, is_causal, cb) \ + [&] { \ + if constexpr (rank == 1) { \ + DISPATCH_FNA_BACKWARD_1D(cc, dtype, is_causal, cb); \ + } \ + else if constexpr (rank == 2) { \ + DISPATCH_FNA_BACKWARD_2D(cc, dtype, is_causal, cb); \ + } \ + else if constexpr (rank == 3) { \ + DISPATCH_FNA_BACKWARD_3D(cc, dtype, is_causal, cb); \ + } \ + else { \ + throw std::runtime_error("NATTEN FNA kernel dispatch failed! Invalid spatial extent rank! Only 1, 2, and 3D are supported!"); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fna + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/fna/kernels.h b/natten/csrc/autogen/include/natten_autogen/cuda/fna/kernels.h new file mode 100644 index 00000000..2fbb3e79 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/fna/kernels.h @@ -0,0 +1,16931 @@ +#pragma once + + +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm50_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm50_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm50_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm50_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm50_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm50_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm50_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm50_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm50_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm50_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm50_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm50_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna1d_sm50_float32_cm_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>(), fna1d_32x128x32_sm50_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>(), fna1d_64x64x32_sm50_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>(), fna1d_64x128x32_sm50_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>(), fna1d_32x128x64_sm50_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>(), fna1d_64x64x64_sm50_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>(), fna1d_64x128x64_sm50_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>(), fna1d_32x128x128_sm50_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>(), fna1d_64x64x128_sm50_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>(), fna1d_64x128x128_sm50_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>(), fna1d_32x128x65536_sm50_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>(), fna1d_64x128x65536_sm50_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna1d_64x64x65536_sm50_float32_cm_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm50_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm50_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm50_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm50_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm50_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm50_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm50_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm50_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm50_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm50_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm50_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm50_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna1d_sm50_float32_cm_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>(), fna1d_32x128x32_sm50_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>(), fna1d_64x64x32_sm50_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>(), fna1d_64x128x32_sm50_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>(), fna1d_32x128x64_sm50_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>(), fna1d_64x64x64_sm50_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>(), fna1d_64x128x64_sm50_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>(), fna1d_32x128x128_sm50_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>(), fna1d_64x64x128_sm50_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>(), fna1d_64x128x128_sm50_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>(), fna1d_32x128x65536_sm50_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>(), fna1d_64x128x65536_sm50_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna1d_64x64x65536_sm50_float32_cm_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm50_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm50_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm50_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm50_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm50_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm50_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm50_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm50_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm50_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm50_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm50_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm50_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna1d_sm50_float16_cm_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>(), fna1d_32x128x32_sm50_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>(), fna1d_64x64x32_sm50_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>(), fna1d_64x128x32_sm50_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>(), fna1d_32x128x64_sm50_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>(), fna1d_64x64x64_sm50_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>(), fna1d_64x128x64_sm50_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>(), fna1d_32x128x128_sm50_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>(), fna1d_64x64x128_sm50_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>(), fna1d_64x128x128_sm50_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>(), fna1d_32x128x65536_sm50_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>(), fna1d_64x128x65536_sm50_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna1d_64x64x65536_sm50_float16_cm_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm50_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm50_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm50_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm50_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm50_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm50_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm50_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm50_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm50_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm50_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm50_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm50_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna1d_sm50_float16_cm_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>(), fna1d_32x128x32_sm50_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>(), fna1d_64x64x32_sm50_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>(), fna1d_64x128x32_sm50_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>(), fna1d_32x128x64_sm50_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>(), fna1d_64x64x64_sm50_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>(), fna1d_64x128x64_sm50_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>(), fna1d_32x128x128_sm50_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>(), fna1d_64x64x128_sm50_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>(), fna1d_64x128x128_sm50_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>(), fna1d_32x128x65536_sm50_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>(), fna1d_64x128x65536_sm50_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna1d_64x64x65536_sm50_float16_cm_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm70_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm70_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm70_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm70_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm70_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm70_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm70_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm70_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm70_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm70_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm70_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm70_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna1d_sm70_float32_cm_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>(), fna1d_32x128x32_sm70_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>(), fna1d_64x64x32_sm70_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>(), fna1d_64x128x32_sm70_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>(), fna1d_32x128x64_sm70_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>(), fna1d_64x64x64_sm70_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>(), fna1d_64x128x64_sm70_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>(), fna1d_32x128x128_sm70_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>(), fna1d_64x64x128_sm70_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>(), fna1d_64x128x128_sm70_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>(), fna1d_32x128x65536_sm70_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>(), fna1d_64x128x65536_sm70_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna1d_64x64x65536_sm70_float32_cm_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm70_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm70_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm70_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm70_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm70_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm70_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm70_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm70_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm70_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm70_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm70_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm70_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna1d_sm70_float32_cm_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>(), fna1d_32x128x32_sm70_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>(), fna1d_64x64x32_sm70_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>(), fna1d_64x128x32_sm70_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>(), fna1d_32x128x64_sm70_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>(), fna1d_64x64x64_sm70_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>(), fna1d_64x128x64_sm70_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>(), fna1d_32x128x128_sm70_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>(), fna1d_64x64x128_sm70_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>(), fna1d_64x128x128_sm70_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>(), fna1d_32x128x65536_sm70_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>(), fna1d_64x128x65536_sm70_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna1d_64x64x65536_sm70_float32_cm_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm70_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm70_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm70_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm70_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm70_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm70_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm70_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm70_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm70_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm70_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm70_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm70_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna1d_sm70_float16_cm_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>(), fna1d_32x128x32_sm70_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>(), fna1d_64x64x32_sm70_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>(), fna1d_64x128x32_sm70_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>(), fna1d_32x128x64_sm70_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>(), fna1d_64x64x64_sm70_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>(), fna1d_64x128x64_sm70_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>(), fna1d_32x128x128_sm70_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>(), fna1d_64x64x128_sm70_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>(), fna1d_64x128x128_sm70_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>(), fna1d_32x128x65536_sm70_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>(), fna1d_64x128x65536_sm70_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna1d_64x64x65536_sm70_float16_cm_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm70_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm70_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm70_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm70_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm70_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm70_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm70_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm70_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm70_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm70_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm70_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm70_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna1d_sm70_float16_cm_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>(), fna1d_32x128x32_sm70_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>(), fna1d_64x64x32_sm70_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>(), fna1d_64x128x32_sm70_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>(), fna1d_32x128x64_sm70_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>(), fna1d_64x64x64_sm70_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>(), fna1d_64x128x64_sm70_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>(), fna1d_32x128x128_sm70_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>(), fna1d_64x64x128_sm70_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>(), fna1d_64x128x128_sm70_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>(), fna1d_32x128x65536_sm70_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>(), fna1d_64x128x65536_sm70_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna1d_64x64x65536_sm70_float16_cm_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm75_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm75_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm75_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm75_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm75_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm75_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm75_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm75_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm75_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm75_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm75_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm75_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna1d_sm75_float32_cm_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>(), fna1d_32x128x32_sm75_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>(), fna1d_64x64x32_sm75_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>(), fna1d_64x128x32_sm75_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>(), fna1d_32x128x64_sm75_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>(), fna1d_64x64x64_sm75_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>(), fna1d_64x128x64_sm75_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>(), fna1d_32x128x128_sm75_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>(), fna1d_64x64x128_sm75_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>(), fna1d_64x128x128_sm75_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>(), fna1d_32x128x65536_sm75_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>(), fna1d_64x128x65536_sm75_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna1d_64x64x65536_sm75_float32_cm_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm75_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm75_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm75_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm75_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm75_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm75_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm75_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm75_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm75_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm75_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm75_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm75_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna1d_sm75_float32_cm_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>(), fna1d_32x128x32_sm75_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>(), fna1d_64x64x32_sm75_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>(), fna1d_64x128x32_sm75_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>(), fna1d_32x128x64_sm75_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>(), fna1d_64x64x64_sm75_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>(), fna1d_64x128x64_sm75_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>(), fna1d_32x128x128_sm75_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>(), fna1d_64x64x128_sm75_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>(), fna1d_64x128x128_sm75_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>(), fna1d_32x128x65536_sm75_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>(), fna1d_64x128x65536_sm75_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna1d_64x64x65536_sm75_float32_cm_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm75_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm75_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm75_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm75_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm75_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm75_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm75_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm75_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm75_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm75_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm75_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm75_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna1d_sm75_float16_cm_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>(), fna1d_32x128x32_sm75_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>(), fna1d_64x64x32_sm75_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>(), fna1d_64x128x32_sm75_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>(), fna1d_32x128x64_sm75_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>(), fna1d_64x64x64_sm75_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>(), fna1d_64x128x64_sm75_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>(), fna1d_32x128x128_sm75_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>(), fna1d_64x64x128_sm75_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>(), fna1d_64x128x128_sm75_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>(), fna1d_32x128x65536_sm75_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>(), fna1d_64x128x65536_sm75_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna1d_64x64x65536_sm75_float16_cm_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm75_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm75_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm75_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm75_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm75_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm75_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm75_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm75_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm75_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm75_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm75_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm75_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna1d_sm75_float16_cm_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>(), fna1d_32x128x32_sm75_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>(), fna1d_64x64x32_sm75_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>(), fna1d_64x128x32_sm75_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>(), fna1d_32x128x64_sm75_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>(), fna1d_64x64x64_sm75_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>(), fna1d_64x128x64_sm75_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>(), fna1d_32x128x128_sm75_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>(), fna1d_64x64x128_sm75_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>(), fna1d_64x128x128_sm75_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>(), fna1d_32x128x65536_sm75_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>(), fna1d_64x128x65536_sm75_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna1d_64x64x65536_sm75_float16_cm_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm80_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm80_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm80_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm80_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm80_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm80_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm80_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm80_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm80_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm80_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm80_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm80_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna1d_sm80_float32_cm_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>(), fna1d_32x128x32_sm80_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>(), fna1d_64x64x32_sm80_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>(), fna1d_64x128x32_sm80_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>(), fna1d_32x128x64_sm80_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>(), fna1d_64x64x64_sm80_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>(), fna1d_64x128x64_sm80_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>(), fna1d_32x128x128_sm80_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>(), fna1d_64x64x128_sm80_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>(), fna1d_64x128x128_sm80_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna1d_32x128x65536_sm80_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna1d_64x128x65536_sm80_float32_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna1d_64x64x65536_sm80_float32_cm_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm80_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm80_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm80_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm80_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm80_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm80_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm80_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm80_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm80_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm80_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm80_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm80_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna1d_sm80_float32_cm_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>(), fna1d_32x128x32_sm80_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>(), fna1d_64x64x32_sm80_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>(), fna1d_64x128x32_sm80_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>(), fna1d_32x128x64_sm80_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>(), fna1d_64x64x64_sm80_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>(), fna1d_64x128x64_sm80_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>(), fna1d_32x128x128_sm80_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>(), fna1d_64x64x128_sm80_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>(), fna1d_64x128x128_sm80_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna1d_32x128x65536_sm80_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna1d_64x128x65536_sm80_float32_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna1d_64x64x65536_sm80_float32_cm_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm80_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm80_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm80_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm80_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm80_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm80_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm80_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm80_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm80_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm80_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm80_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm80_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna1d_sm80_float16_cm_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>(), fna1d_32x128x32_sm80_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna1d_64x64x32_sm80_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>(), fna1d_64x128x32_sm80_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>(), fna1d_32x128x64_sm80_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna1d_64x64x64_sm80_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>(), fna1d_64x128x64_sm80_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>(), fna1d_32x128x128_sm80_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna1d_64x64x128_sm80_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>(), fna1d_64x128x128_sm80_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna1d_32x128x65536_sm80_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna1d_64x128x65536_sm80_float16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna1d_64x64x65536_sm80_float16_cm_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm80_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm80_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm80_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm80_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm80_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm80_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm80_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm80_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm80_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm80_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm80_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm80_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna1d_sm80_float16_cm_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>(), fna1d_32x128x32_sm80_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna1d_64x64x32_sm80_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>(), fna1d_64x128x32_sm80_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>(), fna1d_32x128x64_sm80_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna1d_64x64x64_sm80_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>(), fna1d_64x128x64_sm80_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>(), fna1d_32x128x128_sm80_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna1d_64x64x128_sm80_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>(), fna1d_64x128x128_sm80_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna1d_32x128x65536_sm80_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna1d_64x128x65536_sm80_float16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna1d_64x64x65536_sm80_float16_cm_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna1d_sm80_bfloat16_cm_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>(), fna1d_32x128x32_sm80_bfloat16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna1d_64x64x32_sm80_bfloat16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>(), fna1d_64x128x32_sm80_bfloat16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>(), fna1d_32x128x64_sm80_bfloat16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna1d_64x64x64_sm80_bfloat16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>(), fna1d_64x128x64_sm80_bfloat16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>(), fna1d_32x128x128_sm80_bfloat16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna1d_64x64x128_sm80_bfloat16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>(), fna1d_64x128x128_sm80_bfloat16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna1d_32x128x65536_sm80_bfloat16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna1d_64x128x65536_sm80_bfloat16_cm_0); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna1d_64x64x65536_sm80_bfloat16_cm_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna1d_sm80_bfloat16_cm_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>(), fna1d_32x128x32_sm80_bfloat16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna1d_64x64x32_sm80_bfloat16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>(), fna1d_64x128x32_sm80_bfloat16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>(), fna1d_32x128x64_sm80_bfloat16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna1d_64x64x64_sm80_bfloat16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>(), fna1d_64x128x64_sm80_bfloat16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>(), fna1d_32x128x128_sm80_bfloat16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna1d_64x64x128_sm80_bfloat16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>(), fna1d_64x128x128_sm80_bfloat16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna1d_32x128x65536_sm80_bfloat16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna1d_64x128x65536_sm80_bfloat16_cm_1); + cb(FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna1d_64x64x65536_sm80_bfloat16_cm_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm50_float32_cm_0_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>(), fna2d_32x128x32_sm50_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>(), fna2d_64x64x32_sm50_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>(), fna2d_64x128x32_sm50_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>(), fna2d_32x128x64_sm50_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>(), fna2d_64x64x64_sm50_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>(), fna2d_64x128x64_sm50_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>(), fna2d_32x128x128_sm50_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>(), fna2d_64x64x128_sm50_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>(), fna2d_64x128x128_sm50_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>(), fna2d_32x128x65536_sm50_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>(), fna2d_64x128x65536_sm50_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna2d_64x64x65536_sm50_float32_cm_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm50_float32_cm_0_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>(), fna2d_32x128x32_sm50_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>(), fna2d_64x64x32_sm50_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>(), fna2d_64x128x32_sm50_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>(), fna2d_32x128x64_sm50_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>(), fna2d_64x64x64_sm50_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>(), fna2d_64x128x64_sm50_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>(), fna2d_32x128x128_sm50_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>(), fna2d_64x64x128_sm50_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>(), fna2d_64x128x128_sm50_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>(), fna2d_32x128x65536_sm50_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>(), fna2d_64x128x65536_sm50_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna2d_64x64x65536_sm50_float32_cm_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm50_float32_cm_1_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>(), fna2d_32x128x32_sm50_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>(), fna2d_64x64x32_sm50_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>(), fna2d_64x128x32_sm50_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>(), fna2d_32x128x64_sm50_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>(), fna2d_64x64x64_sm50_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>(), fna2d_64x128x64_sm50_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>(), fna2d_32x128x128_sm50_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>(), fna2d_64x64x128_sm50_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>(), fna2d_64x128x128_sm50_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>(), fna2d_32x128x65536_sm50_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>(), fna2d_64x128x65536_sm50_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna2d_64x64x65536_sm50_float32_cm_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm50_float32_cm_1_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>(), fna2d_32x128x32_sm50_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>(), fna2d_64x64x32_sm50_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>(), fna2d_64x128x32_sm50_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>(), fna2d_32x128x64_sm50_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>(), fna2d_64x64x64_sm50_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>(), fna2d_64x128x64_sm50_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>(), fna2d_32x128x128_sm50_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>(), fna2d_64x64x128_sm50_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>(), fna2d_64x128x128_sm50_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>(), fna2d_32x128x65536_sm50_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>(), fna2d_64x128x65536_sm50_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna2d_64x64x65536_sm50_float32_cm_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm50_float16_cm_0_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>(), fna2d_32x128x32_sm50_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>(), fna2d_64x64x32_sm50_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>(), fna2d_64x128x32_sm50_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>(), fna2d_32x128x64_sm50_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>(), fna2d_64x64x64_sm50_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>(), fna2d_64x128x64_sm50_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>(), fna2d_32x128x128_sm50_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>(), fna2d_64x64x128_sm50_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>(), fna2d_64x128x128_sm50_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>(), fna2d_32x128x65536_sm50_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>(), fna2d_64x128x65536_sm50_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna2d_64x64x65536_sm50_float16_cm_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm50_float16_cm_0_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>(), fna2d_32x128x32_sm50_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>(), fna2d_64x64x32_sm50_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>(), fna2d_64x128x32_sm50_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>(), fna2d_32x128x64_sm50_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>(), fna2d_64x64x64_sm50_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>(), fna2d_64x128x64_sm50_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>(), fna2d_32x128x128_sm50_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>(), fna2d_64x64x128_sm50_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>(), fna2d_64x128x128_sm50_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>(), fna2d_32x128x65536_sm50_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>(), fna2d_64x128x65536_sm50_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna2d_64x64x65536_sm50_float16_cm_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm50_float16_cm_1_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>(), fna2d_32x128x32_sm50_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>(), fna2d_64x64x32_sm50_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>(), fna2d_64x128x32_sm50_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>(), fna2d_32x128x64_sm50_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>(), fna2d_64x64x64_sm50_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>(), fna2d_64x128x64_sm50_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>(), fna2d_32x128x128_sm50_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>(), fna2d_64x64x128_sm50_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>(), fna2d_64x128x128_sm50_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>(), fna2d_32x128x65536_sm50_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>(), fna2d_64x128x65536_sm50_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna2d_64x64x65536_sm50_float16_cm_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm50_float16_cm_1_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>(), fna2d_32x128x32_sm50_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>(), fna2d_64x64x32_sm50_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>(), fna2d_64x128x32_sm50_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>(), fna2d_32x128x64_sm50_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>(), fna2d_64x64x64_sm50_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>(), fna2d_64x128x64_sm50_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>(), fna2d_32x128x128_sm50_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>(), fna2d_64x64x128_sm50_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>(), fna2d_64x128x128_sm50_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>(), fna2d_32x128x65536_sm50_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>(), fna2d_64x128x65536_sm50_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna2d_64x64x65536_sm50_float16_cm_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm70_float32_cm_0_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>(), fna2d_32x128x32_sm70_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>(), fna2d_64x64x32_sm70_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>(), fna2d_64x128x32_sm70_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>(), fna2d_32x128x64_sm70_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>(), fna2d_64x64x64_sm70_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>(), fna2d_64x128x64_sm70_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>(), fna2d_32x128x128_sm70_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>(), fna2d_64x64x128_sm70_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>(), fna2d_64x128x128_sm70_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>(), fna2d_32x128x65536_sm70_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>(), fna2d_64x128x65536_sm70_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna2d_64x64x65536_sm70_float32_cm_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm70_float32_cm_0_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>(), fna2d_32x128x32_sm70_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>(), fna2d_64x64x32_sm70_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>(), fna2d_64x128x32_sm70_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>(), fna2d_32x128x64_sm70_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>(), fna2d_64x64x64_sm70_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>(), fna2d_64x128x64_sm70_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>(), fna2d_32x128x128_sm70_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>(), fna2d_64x64x128_sm70_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>(), fna2d_64x128x128_sm70_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>(), fna2d_32x128x65536_sm70_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>(), fna2d_64x128x65536_sm70_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna2d_64x64x65536_sm70_float32_cm_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm70_float32_cm_1_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>(), fna2d_32x128x32_sm70_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>(), fna2d_64x64x32_sm70_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>(), fna2d_64x128x32_sm70_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>(), fna2d_32x128x64_sm70_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>(), fna2d_64x64x64_sm70_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>(), fna2d_64x128x64_sm70_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>(), fna2d_32x128x128_sm70_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>(), fna2d_64x64x128_sm70_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>(), fna2d_64x128x128_sm70_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>(), fna2d_32x128x65536_sm70_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>(), fna2d_64x128x65536_sm70_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna2d_64x64x65536_sm70_float32_cm_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm70_float32_cm_1_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>(), fna2d_32x128x32_sm70_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>(), fna2d_64x64x32_sm70_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>(), fna2d_64x128x32_sm70_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>(), fna2d_32x128x64_sm70_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>(), fna2d_64x64x64_sm70_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>(), fna2d_64x128x64_sm70_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>(), fna2d_32x128x128_sm70_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>(), fna2d_64x64x128_sm70_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>(), fna2d_64x128x128_sm70_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>(), fna2d_32x128x65536_sm70_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>(), fna2d_64x128x65536_sm70_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna2d_64x64x65536_sm70_float32_cm_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm70_float16_cm_0_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>(), fna2d_32x128x32_sm70_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>(), fna2d_64x64x32_sm70_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>(), fna2d_64x128x32_sm70_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>(), fna2d_32x128x64_sm70_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>(), fna2d_64x64x64_sm70_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>(), fna2d_64x128x64_sm70_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>(), fna2d_32x128x128_sm70_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>(), fna2d_64x64x128_sm70_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>(), fna2d_64x128x128_sm70_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>(), fna2d_32x128x65536_sm70_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>(), fna2d_64x128x65536_sm70_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna2d_64x64x65536_sm70_float16_cm_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm70_float16_cm_0_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>(), fna2d_32x128x32_sm70_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>(), fna2d_64x64x32_sm70_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>(), fna2d_64x128x32_sm70_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>(), fna2d_32x128x64_sm70_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>(), fna2d_64x64x64_sm70_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>(), fna2d_64x128x64_sm70_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>(), fna2d_32x128x128_sm70_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>(), fna2d_64x64x128_sm70_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>(), fna2d_64x128x128_sm70_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>(), fna2d_32x128x65536_sm70_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>(), fna2d_64x128x65536_sm70_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna2d_64x64x65536_sm70_float16_cm_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm70_float16_cm_1_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>(), fna2d_32x128x32_sm70_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>(), fna2d_64x64x32_sm70_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>(), fna2d_64x128x32_sm70_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>(), fna2d_32x128x64_sm70_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>(), fna2d_64x64x64_sm70_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>(), fna2d_64x128x64_sm70_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>(), fna2d_32x128x128_sm70_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>(), fna2d_64x64x128_sm70_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>(), fna2d_64x128x128_sm70_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>(), fna2d_32x128x65536_sm70_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>(), fna2d_64x128x65536_sm70_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna2d_64x64x65536_sm70_float16_cm_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm70_float16_cm_1_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>(), fna2d_32x128x32_sm70_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>(), fna2d_64x64x32_sm70_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>(), fna2d_64x128x32_sm70_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>(), fna2d_32x128x64_sm70_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>(), fna2d_64x64x64_sm70_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>(), fna2d_64x128x64_sm70_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>(), fna2d_32x128x128_sm70_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>(), fna2d_64x64x128_sm70_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>(), fna2d_64x128x128_sm70_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>(), fna2d_32x128x65536_sm70_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>(), fna2d_64x128x65536_sm70_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna2d_64x64x65536_sm70_float16_cm_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm75_float32_cm_0_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>(), fna2d_32x128x32_sm75_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>(), fna2d_64x64x32_sm75_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>(), fna2d_64x128x32_sm75_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>(), fna2d_32x128x64_sm75_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>(), fna2d_64x64x64_sm75_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>(), fna2d_64x128x64_sm75_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>(), fna2d_32x128x128_sm75_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>(), fna2d_64x64x128_sm75_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>(), fna2d_64x128x128_sm75_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>(), fna2d_32x128x65536_sm75_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>(), fna2d_64x128x65536_sm75_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna2d_64x64x65536_sm75_float32_cm_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm75_float32_cm_0_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>(), fna2d_32x128x32_sm75_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>(), fna2d_64x64x32_sm75_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>(), fna2d_64x128x32_sm75_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>(), fna2d_32x128x64_sm75_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>(), fna2d_64x64x64_sm75_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>(), fna2d_64x128x64_sm75_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>(), fna2d_32x128x128_sm75_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>(), fna2d_64x64x128_sm75_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>(), fna2d_64x128x128_sm75_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>(), fna2d_32x128x65536_sm75_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>(), fna2d_64x128x65536_sm75_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna2d_64x64x65536_sm75_float32_cm_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm75_float32_cm_1_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>(), fna2d_32x128x32_sm75_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>(), fna2d_64x64x32_sm75_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>(), fna2d_64x128x32_sm75_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>(), fna2d_32x128x64_sm75_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>(), fna2d_64x64x64_sm75_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>(), fna2d_64x128x64_sm75_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>(), fna2d_32x128x128_sm75_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>(), fna2d_64x64x128_sm75_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>(), fna2d_64x128x128_sm75_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>(), fna2d_32x128x65536_sm75_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>(), fna2d_64x128x65536_sm75_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna2d_64x64x65536_sm75_float32_cm_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm75_float32_cm_1_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>(), fna2d_32x128x32_sm75_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>(), fna2d_64x64x32_sm75_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>(), fna2d_64x128x32_sm75_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>(), fna2d_32x128x64_sm75_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>(), fna2d_64x64x64_sm75_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>(), fna2d_64x128x64_sm75_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>(), fna2d_32x128x128_sm75_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>(), fna2d_64x64x128_sm75_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>(), fna2d_64x128x128_sm75_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>(), fna2d_32x128x65536_sm75_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>(), fna2d_64x128x65536_sm75_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna2d_64x64x65536_sm75_float32_cm_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm75_float16_cm_0_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>(), fna2d_32x128x32_sm75_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>(), fna2d_64x64x32_sm75_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>(), fna2d_64x128x32_sm75_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>(), fna2d_32x128x64_sm75_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>(), fna2d_64x64x64_sm75_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>(), fna2d_64x128x64_sm75_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>(), fna2d_32x128x128_sm75_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>(), fna2d_64x64x128_sm75_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>(), fna2d_64x128x128_sm75_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>(), fna2d_32x128x65536_sm75_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>(), fna2d_64x128x65536_sm75_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna2d_64x64x65536_sm75_float16_cm_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm75_float16_cm_0_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>(), fna2d_32x128x32_sm75_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>(), fna2d_64x64x32_sm75_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>(), fna2d_64x128x32_sm75_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>(), fna2d_32x128x64_sm75_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>(), fna2d_64x64x64_sm75_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>(), fna2d_64x128x64_sm75_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>(), fna2d_32x128x128_sm75_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>(), fna2d_64x64x128_sm75_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>(), fna2d_64x128x128_sm75_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>(), fna2d_32x128x65536_sm75_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>(), fna2d_64x128x65536_sm75_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna2d_64x64x65536_sm75_float16_cm_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm75_float16_cm_1_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>(), fna2d_32x128x32_sm75_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>(), fna2d_64x64x32_sm75_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>(), fna2d_64x128x32_sm75_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>(), fna2d_32x128x64_sm75_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>(), fna2d_64x64x64_sm75_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>(), fna2d_64x128x64_sm75_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>(), fna2d_32x128x128_sm75_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>(), fna2d_64x64x128_sm75_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>(), fna2d_64x128x128_sm75_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>(), fna2d_32x128x65536_sm75_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>(), fna2d_64x128x65536_sm75_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna2d_64x64x65536_sm75_float16_cm_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm75_float16_cm_1_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>(), fna2d_32x128x32_sm75_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>(), fna2d_64x64x32_sm75_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>(), fna2d_64x128x32_sm75_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>(), fna2d_32x128x64_sm75_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>(), fna2d_64x64x64_sm75_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>(), fna2d_64x128x64_sm75_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>(), fna2d_32x128x128_sm75_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>(), fna2d_64x64x128_sm75_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>(), fna2d_64x128x128_sm75_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>(), fna2d_32x128x65536_sm75_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>(), fna2d_64x128x65536_sm75_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna2d_64x64x65536_sm75_float16_cm_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm80_float32_cm_0_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>(), fna2d_32x128x32_sm80_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>(), fna2d_64x64x32_sm80_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>(), fna2d_64x128x32_sm80_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>(), fna2d_32x128x64_sm80_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>(), fna2d_64x64x64_sm80_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>(), fna2d_64x128x64_sm80_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>(), fna2d_32x128x128_sm80_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>(), fna2d_64x64x128_sm80_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>(), fna2d_64x128x128_sm80_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna2d_32x128x65536_sm80_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna2d_64x128x65536_sm80_float32_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna2d_64x64x65536_sm80_float32_cm_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm80_float32_cm_0_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>(), fna2d_32x128x32_sm80_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>(), fna2d_64x64x32_sm80_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>(), fna2d_64x128x32_sm80_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>(), fna2d_32x128x64_sm80_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>(), fna2d_64x64x64_sm80_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>(), fna2d_64x128x64_sm80_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>(), fna2d_32x128x128_sm80_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>(), fna2d_64x64x128_sm80_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>(), fna2d_64x128x128_sm80_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna2d_32x128x65536_sm80_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna2d_64x128x65536_sm80_float32_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna2d_64x64x65536_sm80_float32_cm_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm80_float32_cm_1_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>(), fna2d_32x128x32_sm80_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>(), fna2d_64x64x32_sm80_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>(), fna2d_64x128x32_sm80_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>(), fna2d_32x128x64_sm80_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>(), fna2d_64x64x64_sm80_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>(), fna2d_64x128x64_sm80_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>(), fna2d_32x128x128_sm80_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>(), fna2d_64x64x128_sm80_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>(), fna2d_64x128x128_sm80_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna2d_32x128x65536_sm80_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna2d_64x128x65536_sm80_float32_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna2d_64x64x65536_sm80_float32_cm_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm80_float32_cm_1_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>(), fna2d_32x128x32_sm80_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>(), fna2d_64x64x32_sm80_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>(), fna2d_64x128x32_sm80_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>(), fna2d_32x128x64_sm80_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>(), fna2d_64x64x64_sm80_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>(), fna2d_64x128x64_sm80_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>(), fna2d_32x128x128_sm80_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>(), fna2d_64x64x128_sm80_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>(), fna2d_64x128x128_sm80_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna2d_32x128x65536_sm80_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna2d_64x128x65536_sm80_float32_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna2d_64x64x65536_sm80_float32_cm_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm80_float16_cm_0_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>(), fna2d_32x128x32_sm80_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna2d_64x64x32_sm80_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>(), fna2d_64x128x32_sm80_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>(), fna2d_32x128x64_sm80_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna2d_64x64x64_sm80_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>(), fna2d_64x128x64_sm80_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>(), fna2d_32x128x128_sm80_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna2d_64x64x128_sm80_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>(), fna2d_64x128x128_sm80_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna2d_32x128x65536_sm80_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna2d_64x128x65536_sm80_float16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna2d_64x64x65536_sm80_float16_cm_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm80_float16_cm_0_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>(), fna2d_32x128x32_sm80_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna2d_64x64x32_sm80_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>(), fna2d_64x128x32_sm80_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>(), fna2d_32x128x64_sm80_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna2d_64x64x64_sm80_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>(), fna2d_64x128x64_sm80_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>(), fna2d_32x128x128_sm80_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna2d_64x64x128_sm80_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>(), fna2d_64x128x128_sm80_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna2d_32x128x65536_sm80_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna2d_64x128x65536_sm80_float16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna2d_64x64x65536_sm80_float16_cm_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm80_float16_cm_1_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>(), fna2d_32x128x32_sm80_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna2d_64x64x32_sm80_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>(), fna2d_64x128x32_sm80_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>(), fna2d_32x128x64_sm80_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna2d_64x64x64_sm80_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>(), fna2d_64x128x64_sm80_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>(), fna2d_32x128x128_sm80_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna2d_64x64x128_sm80_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>(), fna2d_64x128x128_sm80_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna2d_32x128x65536_sm80_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna2d_64x128x65536_sm80_float16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna2d_64x64x65536_sm80_float16_cm_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm80_float16_cm_1_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>(), fna2d_32x128x32_sm80_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna2d_64x64x32_sm80_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>(), fna2d_64x128x32_sm80_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>(), fna2d_32x128x64_sm80_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna2d_64x64x64_sm80_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>(), fna2d_64x128x64_sm80_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>(), fna2d_32x128x128_sm80_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna2d_64x64x128_sm80_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>(), fna2d_64x128x128_sm80_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna2d_32x128x65536_sm80_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna2d_64x128x65536_sm80_float16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna2d_64x64x65536_sm80_float16_cm_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm80_bfloat16_cm_0_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>(), fna2d_32x128x32_sm80_bfloat16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna2d_64x64x32_sm80_bfloat16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>(), fna2d_64x128x32_sm80_bfloat16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>(), fna2d_32x128x64_sm80_bfloat16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna2d_64x64x64_sm80_bfloat16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>(), fna2d_64x128x64_sm80_bfloat16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>(), fna2d_32x128x128_sm80_bfloat16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna2d_64x64x128_sm80_bfloat16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>(), fna2d_64x128x128_sm80_bfloat16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna2d_32x128x65536_sm80_bfloat16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna2d_64x128x65536_sm80_bfloat16_cm_0_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna2d_64x64x65536_sm80_bfloat16_cm_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm80_bfloat16_cm_0_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>(), fna2d_32x128x32_sm80_bfloat16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna2d_64x64x32_sm80_bfloat16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>(), fna2d_64x128x32_sm80_bfloat16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>(), fna2d_32x128x64_sm80_bfloat16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna2d_64x64x64_sm80_bfloat16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>(), fna2d_64x128x64_sm80_bfloat16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>(), fna2d_32x128x128_sm80_bfloat16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna2d_64x64x128_sm80_bfloat16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>(), fna2d_64x128x128_sm80_bfloat16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna2d_32x128x65536_sm80_bfloat16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna2d_64x128x65536_sm80_bfloat16_cm_0_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna2d_64x64x65536_sm80_bfloat16_cm_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm80_bfloat16_cm_1_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>(), fna2d_32x128x32_sm80_bfloat16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna2d_64x64x32_sm80_bfloat16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>(), fna2d_64x128x32_sm80_bfloat16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>(), fna2d_32x128x64_sm80_bfloat16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna2d_64x64x64_sm80_bfloat16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>(), fna2d_64x128x64_sm80_bfloat16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>(), fna2d_32x128x128_sm80_bfloat16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna2d_64x64x128_sm80_bfloat16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>(), fna2d_64x128x128_sm80_bfloat16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna2d_32x128x65536_sm80_bfloat16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna2d_64x128x65536_sm80_bfloat16_cm_1_0); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna2d_64x64x65536_sm80_bfloat16_cm_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna2d_sm80_bfloat16_cm_1_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>(), fna2d_32x128x32_sm80_bfloat16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna2d_64x64x32_sm80_bfloat16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>(), fna2d_64x128x32_sm80_bfloat16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>(), fna2d_32x128x64_sm80_bfloat16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna2d_64x64x64_sm80_bfloat16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>(), fna2d_64x128x64_sm80_bfloat16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>(), fna2d_32x128x128_sm80_bfloat16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna2d_64x64x128_sm80_bfloat16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>(), fna2d_64x128x128_sm80_bfloat16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna2d_32x128x65536_sm80_bfloat16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna2d_64x128x65536_sm80_bfloat16_cm_1_1); + cb(FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna2d_64x64x65536_sm80_bfloat16_cm_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm50_float32_cm_0_0_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>(), fna3d_32x128x32_sm50_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_64x64x32_sm50_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>(), fna3d_64x128x32_sm50_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>(), fna3d_32x128x64_sm50_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_64x64x64_sm50_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>(), fna3d_64x128x64_sm50_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>(), fna3d_32x128x128_sm50_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_64x64x128_sm50_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>(), fna3d_64x128x128_sm50_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>(), fna3d_32x128x65536_sm50_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>(), fna3d_64x128x65536_sm50_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_64x64x65536_sm50_float32_cm_0_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm50_float32_cm_0_0_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>(), fna3d_32x128x32_sm50_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_64x64x32_sm50_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>(), fna3d_64x128x32_sm50_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>(), fna3d_32x128x64_sm50_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_64x64x64_sm50_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>(), fna3d_64x128x64_sm50_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>(), fna3d_32x128x128_sm50_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_64x64x128_sm50_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>(), fna3d_64x128x128_sm50_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>(), fna3d_32x128x65536_sm50_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>(), fna3d_64x128x65536_sm50_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_64x64x65536_sm50_float32_cm_0_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm50_float32_cm_0_1_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>(), fna3d_32x128x32_sm50_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_64x64x32_sm50_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>(), fna3d_64x128x32_sm50_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>(), fna3d_32x128x64_sm50_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_64x64x64_sm50_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>(), fna3d_64x128x64_sm50_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>(), fna3d_32x128x128_sm50_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_64x64x128_sm50_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>(), fna3d_64x128x128_sm50_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>(), fna3d_32x128x65536_sm50_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>(), fna3d_64x128x65536_sm50_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_64x64x65536_sm50_float32_cm_0_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm50_float32_cm_0_1_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>(), fna3d_32x128x32_sm50_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_64x64x32_sm50_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>(), fna3d_64x128x32_sm50_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>(), fna3d_32x128x64_sm50_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_64x64x64_sm50_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>(), fna3d_64x128x64_sm50_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>(), fna3d_32x128x128_sm50_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_64x64x128_sm50_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>(), fna3d_64x128x128_sm50_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>(), fna3d_32x128x65536_sm50_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>(), fna3d_64x128x65536_sm50_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_64x64x65536_sm50_float32_cm_0_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm50_float32_cm_1_0_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>(), fna3d_32x128x32_sm50_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_64x64x32_sm50_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>(), fna3d_64x128x32_sm50_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>(), fna3d_32x128x64_sm50_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_64x64x64_sm50_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>(), fna3d_64x128x64_sm50_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>(), fna3d_32x128x128_sm50_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_64x64x128_sm50_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>(), fna3d_64x128x128_sm50_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>(), fna3d_32x128x65536_sm50_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>(), fna3d_64x128x65536_sm50_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_64x64x65536_sm50_float32_cm_1_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm50_float32_cm_1_0_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>(), fna3d_32x128x32_sm50_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_64x64x32_sm50_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>(), fna3d_64x128x32_sm50_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>(), fna3d_32x128x64_sm50_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_64x64x64_sm50_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>(), fna3d_64x128x64_sm50_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>(), fna3d_32x128x128_sm50_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_64x64x128_sm50_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>(), fna3d_64x128x128_sm50_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>(), fna3d_32x128x65536_sm50_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>(), fna3d_64x128x65536_sm50_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_64x64x65536_sm50_float32_cm_1_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm50_float32_cm_1_1_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>(), fna3d_32x128x32_sm50_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_64x64x32_sm50_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>(), fna3d_64x128x32_sm50_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>(), fna3d_32x128x64_sm50_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_64x64x64_sm50_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>(), fna3d_64x128x64_sm50_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>(), fna3d_32x128x128_sm50_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_64x64x128_sm50_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>(), fna3d_64x128x128_sm50_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>(), fna3d_32x128x65536_sm50_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>(), fna3d_64x128x65536_sm50_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_64x64x65536_sm50_float32_cm_1_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm50_float32_cm_1_1_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>(), fna3d_32x128x32_sm50_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_64x64x32_sm50_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>(), fna3d_64x128x32_sm50_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>(), fna3d_32x128x64_sm50_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_64x64x64_sm50_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>(), fna3d_64x128x64_sm50_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>(), fna3d_32x128x128_sm50_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_64x64x128_sm50_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>(), fna3d_64x128x128_sm50_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>(), fna3d_32x128x65536_sm50_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>(), fna3d_64x128x65536_sm50_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_64x64x65536_sm50_float32_cm_1_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm50_float16_cm_0_0_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>(), fna3d_32x128x32_sm50_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_64x64x32_sm50_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>(), fna3d_64x128x32_sm50_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>(), fna3d_32x128x64_sm50_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_64x64x64_sm50_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>(), fna3d_64x128x64_sm50_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>(), fna3d_32x128x128_sm50_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_64x64x128_sm50_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>(), fna3d_64x128x128_sm50_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>(), fna3d_32x128x65536_sm50_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>(), fna3d_64x128x65536_sm50_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_64x64x65536_sm50_float16_cm_0_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm50_float16_cm_0_0_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>(), fna3d_32x128x32_sm50_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_64x64x32_sm50_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>(), fna3d_64x128x32_sm50_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>(), fna3d_32x128x64_sm50_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_64x64x64_sm50_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>(), fna3d_64x128x64_sm50_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>(), fna3d_32x128x128_sm50_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_64x64x128_sm50_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>(), fna3d_64x128x128_sm50_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>(), fna3d_32x128x65536_sm50_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>(), fna3d_64x128x65536_sm50_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_64x64x65536_sm50_float16_cm_0_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm50_float16_cm_0_1_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>(), fna3d_32x128x32_sm50_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_64x64x32_sm50_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>(), fna3d_64x128x32_sm50_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>(), fna3d_32x128x64_sm50_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_64x64x64_sm50_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>(), fna3d_64x128x64_sm50_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>(), fna3d_32x128x128_sm50_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_64x64x128_sm50_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>(), fna3d_64x128x128_sm50_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>(), fna3d_32x128x65536_sm50_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>(), fna3d_64x128x65536_sm50_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_64x64x65536_sm50_float16_cm_0_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm50_float16_cm_0_1_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>(), fna3d_32x128x32_sm50_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_64x64x32_sm50_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>(), fna3d_64x128x32_sm50_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>(), fna3d_32x128x64_sm50_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_64x64x64_sm50_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>(), fna3d_64x128x64_sm50_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>(), fna3d_32x128x128_sm50_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_64x64x128_sm50_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>(), fna3d_64x128x128_sm50_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>(), fna3d_32x128x65536_sm50_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>(), fna3d_64x128x65536_sm50_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_64x64x65536_sm50_float16_cm_0_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm50_float16_cm_1_0_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>(), fna3d_32x128x32_sm50_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_64x64x32_sm50_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>(), fna3d_64x128x32_sm50_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>(), fna3d_32x128x64_sm50_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_64x64x64_sm50_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>(), fna3d_64x128x64_sm50_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>(), fna3d_32x128x128_sm50_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_64x64x128_sm50_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>(), fna3d_64x128x128_sm50_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>(), fna3d_32x128x65536_sm50_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>(), fna3d_64x128x65536_sm50_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_64x64x65536_sm50_float16_cm_1_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm50_float16_cm_1_0_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>(), fna3d_32x128x32_sm50_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_64x64x32_sm50_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>(), fna3d_64x128x32_sm50_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>(), fna3d_32x128x64_sm50_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_64x64x64_sm50_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>(), fna3d_64x128x64_sm50_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>(), fna3d_32x128x128_sm50_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_64x64x128_sm50_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>(), fna3d_64x128x128_sm50_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>(), fna3d_32x128x65536_sm50_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>(), fna3d_64x128x65536_sm50_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_64x64x65536_sm50_float16_cm_1_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm50_float16_cm_1_1_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>(), fna3d_32x128x32_sm50_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_64x64x32_sm50_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>(), fna3d_64x128x32_sm50_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>(), fna3d_32x128x64_sm50_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_64x64x64_sm50_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>(), fna3d_64x128x64_sm50_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>(), fna3d_32x128x128_sm50_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_64x64x128_sm50_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>(), fna3d_64x128x128_sm50_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>(), fna3d_32x128x65536_sm50_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>(), fna3d_64x128x65536_sm50_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_64x64x65536_sm50_float16_cm_1_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm50_float16_cm_1_1_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>(), fna3d_32x128x32_sm50_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_64x64x32_sm50_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>(), fna3d_64x128x32_sm50_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>(), fna3d_32x128x64_sm50_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_64x64x64_sm50_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>(), fna3d_64x128x64_sm50_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>(), fna3d_32x128x128_sm50_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_64x64x128_sm50_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>(), fna3d_64x128x128_sm50_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>(), fna3d_32x128x65536_sm50_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>(), fna3d_64x128x65536_sm50_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_64x64x65536_sm50_float16_cm_1_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm70_float32_cm_0_0_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>(), fna3d_32x128x32_sm70_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_64x64x32_sm70_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>(), fna3d_64x128x32_sm70_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>(), fna3d_32x128x64_sm70_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_64x64x64_sm70_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>(), fna3d_64x128x64_sm70_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>(), fna3d_32x128x128_sm70_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_64x64x128_sm70_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>(), fna3d_64x128x128_sm70_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>(), fna3d_32x128x65536_sm70_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>(), fna3d_64x128x65536_sm70_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_64x64x65536_sm70_float32_cm_0_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm70_float32_cm_0_0_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>(), fna3d_32x128x32_sm70_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_64x64x32_sm70_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>(), fna3d_64x128x32_sm70_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>(), fna3d_32x128x64_sm70_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_64x64x64_sm70_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>(), fna3d_64x128x64_sm70_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>(), fna3d_32x128x128_sm70_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_64x64x128_sm70_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>(), fna3d_64x128x128_sm70_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>(), fna3d_32x128x65536_sm70_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>(), fna3d_64x128x65536_sm70_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_64x64x65536_sm70_float32_cm_0_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm70_float32_cm_0_1_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>(), fna3d_32x128x32_sm70_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_64x64x32_sm70_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>(), fna3d_64x128x32_sm70_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>(), fna3d_32x128x64_sm70_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_64x64x64_sm70_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>(), fna3d_64x128x64_sm70_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>(), fna3d_32x128x128_sm70_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_64x64x128_sm70_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>(), fna3d_64x128x128_sm70_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>(), fna3d_32x128x65536_sm70_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>(), fna3d_64x128x65536_sm70_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_64x64x65536_sm70_float32_cm_0_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm70_float32_cm_0_1_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>(), fna3d_32x128x32_sm70_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_64x64x32_sm70_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>(), fna3d_64x128x32_sm70_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>(), fna3d_32x128x64_sm70_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_64x64x64_sm70_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>(), fna3d_64x128x64_sm70_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>(), fna3d_32x128x128_sm70_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_64x64x128_sm70_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>(), fna3d_64x128x128_sm70_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>(), fna3d_32x128x65536_sm70_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>(), fna3d_64x128x65536_sm70_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_64x64x65536_sm70_float32_cm_0_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm70_float32_cm_1_0_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>(), fna3d_32x128x32_sm70_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_64x64x32_sm70_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>(), fna3d_64x128x32_sm70_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>(), fna3d_32x128x64_sm70_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_64x64x64_sm70_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>(), fna3d_64x128x64_sm70_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>(), fna3d_32x128x128_sm70_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_64x64x128_sm70_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>(), fna3d_64x128x128_sm70_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>(), fna3d_32x128x65536_sm70_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>(), fna3d_64x128x65536_sm70_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_64x64x65536_sm70_float32_cm_1_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm70_float32_cm_1_0_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>(), fna3d_32x128x32_sm70_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_64x64x32_sm70_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>(), fna3d_64x128x32_sm70_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>(), fna3d_32x128x64_sm70_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_64x64x64_sm70_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>(), fna3d_64x128x64_sm70_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>(), fna3d_32x128x128_sm70_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_64x64x128_sm70_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>(), fna3d_64x128x128_sm70_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>(), fna3d_32x128x65536_sm70_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>(), fna3d_64x128x65536_sm70_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_64x64x65536_sm70_float32_cm_1_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm70_float32_cm_1_1_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>(), fna3d_32x128x32_sm70_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_64x64x32_sm70_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>(), fna3d_64x128x32_sm70_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>(), fna3d_32x128x64_sm70_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_64x64x64_sm70_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>(), fna3d_64x128x64_sm70_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>(), fna3d_32x128x128_sm70_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_64x64x128_sm70_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>(), fna3d_64x128x128_sm70_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>(), fna3d_32x128x65536_sm70_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>(), fna3d_64x128x65536_sm70_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_64x64x65536_sm70_float32_cm_1_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm70_float32_cm_1_1_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>(), fna3d_32x128x32_sm70_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_64x64x32_sm70_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>(), fna3d_64x128x32_sm70_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>(), fna3d_32x128x64_sm70_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_64x64x64_sm70_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>(), fna3d_64x128x64_sm70_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>(), fna3d_32x128x128_sm70_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_64x64x128_sm70_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>(), fna3d_64x128x128_sm70_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>(), fna3d_32x128x65536_sm70_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>(), fna3d_64x128x65536_sm70_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_64x64x65536_sm70_float32_cm_1_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm70_float16_cm_0_0_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>(), fna3d_32x128x32_sm70_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_64x64x32_sm70_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>(), fna3d_64x128x32_sm70_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>(), fna3d_32x128x64_sm70_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_64x64x64_sm70_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>(), fna3d_64x128x64_sm70_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>(), fna3d_32x128x128_sm70_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_64x64x128_sm70_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>(), fna3d_64x128x128_sm70_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>(), fna3d_32x128x65536_sm70_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>(), fna3d_64x128x65536_sm70_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_64x64x65536_sm70_float16_cm_0_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm70_float16_cm_0_0_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>(), fna3d_32x128x32_sm70_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_64x64x32_sm70_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>(), fna3d_64x128x32_sm70_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>(), fna3d_32x128x64_sm70_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_64x64x64_sm70_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>(), fna3d_64x128x64_sm70_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>(), fna3d_32x128x128_sm70_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_64x64x128_sm70_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>(), fna3d_64x128x128_sm70_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>(), fna3d_32x128x65536_sm70_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>(), fna3d_64x128x65536_sm70_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_64x64x65536_sm70_float16_cm_0_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm70_float16_cm_0_1_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>(), fna3d_32x128x32_sm70_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_64x64x32_sm70_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>(), fna3d_64x128x32_sm70_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>(), fna3d_32x128x64_sm70_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_64x64x64_sm70_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>(), fna3d_64x128x64_sm70_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>(), fna3d_32x128x128_sm70_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_64x64x128_sm70_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>(), fna3d_64x128x128_sm70_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>(), fna3d_32x128x65536_sm70_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>(), fna3d_64x128x65536_sm70_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_64x64x65536_sm70_float16_cm_0_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm70_float16_cm_0_1_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>(), fna3d_32x128x32_sm70_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_64x64x32_sm70_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>(), fna3d_64x128x32_sm70_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>(), fna3d_32x128x64_sm70_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_64x64x64_sm70_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>(), fna3d_64x128x64_sm70_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>(), fna3d_32x128x128_sm70_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_64x64x128_sm70_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>(), fna3d_64x128x128_sm70_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>(), fna3d_32x128x65536_sm70_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>(), fna3d_64x128x65536_sm70_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_64x64x65536_sm70_float16_cm_0_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm70_float16_cm_1_0_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>(), fna3d_32x128x32_sm70_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_64x64x32_sm70_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>(), fna3d_64x128x32_sm70_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>(), fna3d_32x128x64_sm70_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_64x64x64_sm70_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>(), fna3d_64x128x64_sm70_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>(), fna3d_32x128x128_sm70_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_64x64x128_sm70_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>(), fna3d_64x128x128_sm70_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>(), fna3d_32x128x65536_sm70_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>(), fna3d_64x128x65536_sm70_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_64x64x65536_sm70_float16_cm_1_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm70_float16_cm_1_0_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>(), fna3d_32x128x32_sm70_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_64x64x32_sm70_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>(), fna3d_64x128x32_sm70_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>(), fna3d_32x128x64_sm70_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_64x64x64_sm70_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>(), fna3d_64x128x64_sm70_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>(), fna3d_32x128x128_sm70_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_64x64x128_sm70_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>(), fna3d_64x128x128_sm70_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>(), fna3d_32x128x65536_sm70_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>(), fna3d_64x128x65536_sm70_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_64x64x65536_sm70_float16_cm_1_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm70_float16_cm_1_1_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>(), fna3d_32x128x32_sm70_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_64x64x32_sm70_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>(), fna3d_64x128x32_sm70_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>(), fna3d_32x128x64_sm70_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_64x64x64_sm70_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>(), fna3d_64x128x64_sm70_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>(), fna3d_32x128x128_sm70_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_64x64x128_sm70_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>(), fna3d_64x128x128_sm70_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>(), fna3d_32x128x65536_sm70_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>(), fna3d_64x128x65536_sm70_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_64x64x65536_sm70_float16_cm_1_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm70_float16_cm_1_1_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>(), fna3d_32x128x32_sm70_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_64x64x32_sm70_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>(), fna3d_64x128x32_sm70_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>(), fna3d_32x128x64_sm70_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_64x64x64_sm70_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>(), fna3d_64x128x64_sm70_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>(), fna3d_32x128x128_sm70_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_64x64x128_sm70_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>(), fna3d_64x128x128_sm70_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>(), fna3d_32x128x65536_sm70_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>(), fna3d_64x128x65536_sm70_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_64x64x65536_sm70_float16_cm_1_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm75_float32_cm_0_0_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>(), fna3d_32x128x32_sm75_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_64x64x32_sm75_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>(), fna3d_64x128x32_sm75_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>(), fna3d_32x128x64_sm75_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_64x64x64_sm75_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>(), fna3d_64x128x64_sm75_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>(), fna3d_32x128x128_sm75_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_64x64x128_sm75_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>(), fna3d_64x128x128_sm75_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>(), fna3d_32x128x65536_sm75_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>(), fna3d_64x128x65536_sm75_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_64x64x65536_sm75_float32_cm_0_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm75_float32_cm_0_0_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>(), fna3d_32x128x32_sm75_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_64x64x32_sm75_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>(), fna3d_64x128x32_sm75_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>(), fna3d_32x128x64_sm75_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_64x64x64_sm75_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>(), fna3d_64x128x64_sm75_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>(), fna3d_32x128x128_sm75_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_64x64x128_sm75_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>(), fna3d_64x128x128_sm75_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>(), fna3d_32x128x65536_sm75_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>(), fna3d_64x128x65536_sm75_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_64x64x65536_sm75_float32_cm_0_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm75_float32_cm_0_1_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>(), fna3d_32x128x32_sm75_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_64x64x32_sm75_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>(), fna3d_64x128x32_sm75_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>(), fna3d_32x128x64_sm75_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_64x64x64_sm75_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>(), fna3d_64x128x64_sm75_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>(), fna3d_32x128x128_sm75_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_64x64x128_sm75_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>(), fna3d_64x128x128_sm75_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>(), fna3d_32x128x65536_sm75_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>(), fna3d_64x128x65536_sm75_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_64x64x65536_sm75_float32_cm_0_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm75_float32_cm_0_1_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>(), fna3d_32x128x32_sm75_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_64x64x32_sm75_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>(), fna3d_64x128x32_sm75_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>(), fna3d_32x128x64_sm75_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_64x64x64_sm75_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>(), fna3d_64x128x64_sm75_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>(), fna3d_32x128x128_sm75_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_64x64x128_sm75_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>(), fna3d_64x128x128_sm75_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>(), fna3d_32x128x65536_sm75_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>(), fna3d_64x128x65536_sm75_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_64x64x65536_sm75_float32_cm_0_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm75_float32_cm_1_0_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>(), fna3d_32x128x32_sm75_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_64x64x32_sm75_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>(), fna3d_64x128x32_sm75_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>(), fna3d_32x128x64_sm75_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_64x64x64_sm75_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>(), fna3d_64x128x64_sm75_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>(), fna3d_32x128x128_sm75_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_64x64x128_sm75_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>(), fna3d_64x128x128_sm75_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>(), fna3d_32x128x65536_sm75_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>(), fna3d_64x128x65536_sm75_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_64x64x65536_sm75_float32_cm_1_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm75_float32_cm_1_0_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>(), fna3d_32x128x32_sm75_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_64x64x32_sm75_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>(), fna3d_64x128x32_sm75_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>(), fna3d_32x128x64_sm75_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_64x64x64_sm75_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>(), fna3d_64x128x64_sm75_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>(), fna3d_32x128x128_sm75_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_64x64x128_sm75_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>(), fna3d_64x128x128_sm75_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>(), fna3d_32x128x65536_sm75_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>(), fna3d_64x128x65536_sm75_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_64x64x65536_sm75_float32_cm_1_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm75_float32_cm_1_1_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>(), fna3d_32x128x32_sm75_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_64x64x32_sm75_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>(), fna3d_64x128x32_sm75_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>(), fna3d_32x128x64_sm75_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_64x64x64_sm75_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>(), fna3d_64x128x64_sm75_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>(), fna3d_32x128x128_sm75_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_64x64x128_sm75_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>(), fna3d_64x128x128_sm75_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>(), fna3d_32x128x65536_sm75_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>(), fna3d_64x128x65536_sm75_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_64x64x65536_sm75_float32_cm_1_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm75_float32_cm_1_1_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>(), fna3d_32x128x32_sm75_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_64x64x32_sm75_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>(), fna3d_64x128x32_sm75_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>(), fna3d_32x128x64_sm75_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_64x64x64_sm75_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>(), fna3d_64x128x64_sm75_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>(), fna3d_32x128x128_sm75_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_64x64x128_sm75_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>(), fna3d_64x128x128_sm75_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>(), fna3d_32x128x65536_sm75_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>(), fna3d_64x128x65536_sm75_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_64x64x65536_sm75_float32_cm_1_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm75_float16_cm_0_0_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>(), fna3d_32x128x32_sm75_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_64x64x32_sm75_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>(), fna3d_64x128x32_sm75_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>(), fna3d_32x128x64_sm75_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_64x64x64_sm75_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>(), fna3d_64x128x64_sm75_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>(), fna3d_32x128x128_sm75_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_64x64x128_sm75_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>(), fna3d_64x128x128_sm75_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>(), fna3d_32x128x65536_sm75_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>(), fna3d_64x128x65536_sm75_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_64x64x65536_sm75_float16_cm_0_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm75_float16_cm_0_0_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>(), fna3d_32x128x32_sm75_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_64x64x32_sm75_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>(), fna3d_64x128x32_sm75_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>(), fna3d_32x128x64_sm75_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_64x64x64_sm75_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>(), fna3d_64x128x64_sm75_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>(), fna3d_32x128x128_sm75_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_64x64x128_sm75_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>(), fna3d_64x128x128_sm75_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>(), fna3d_32x128x65536_sm75_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>(), fna3d_64x128x65536_sm75_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_64x64x65536_sm75_float16_cm_0_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm75_float16_cm_0_1_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>(), fna3d_32x128x32_sm75_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_64x64x32_sm75_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>(), fna3d_64x128x32_sm75_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>(), fna3d_32x128x64_sm75_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_64x64x64_sm75_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>(), fna3d_64x128x64_sm75_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>(), fna3d_32x128x128_sm75_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_64x64x128_sm75_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>(), fna3d_64x128x128_sm75_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>(), fna3d_32x128x65536_sm75_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>(), fna3d_64x128x65536_sm75_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_64x64x65536_sm75_float16_cm_0_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm75_float16_cm_0_1_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>(), fna3d_32x128x32_sm75_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_64x64x32_sm75_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>(), fna3d_64x128x32_sm75_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>(), fna3d_32x128x64_sm75_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_64x64x64_sm75_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>(), fna3d_64x128x64_sm75_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>(), fna3d_32x128x128_sm75_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_64x64x128_sm75_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>(), fna3d_64x128x128_sm75_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>(), fna3d_32x128x65536_sm75_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>(), fna3d_64x128x65536_sm75_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_64x64x65536_sm75_float16_cm_0_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm75_float16_cm_1_0_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>(), fna3d_32x128x32_sm75_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_64x64x32_sm75_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>(), fna3d_64x128x32_sm75_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>(), fna3d_32x128x64_sm75_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_64x64x64_sm75_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>(), fna3d_64x128x64_sm75_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>(), fna3d_32x128x128_sm75_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_64x64x128_sm75_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>(), fna3d_64x128x128_sm75_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>(), fna3d_32x128x65536_sm75_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>(), fna3d_64x128x65536_sm75_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_64x64x65536_sm75_float16_cm_1_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm75_float16_cm_1_0_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>(), fna3d_32x128x32_sm75_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_64x64x32_sm75_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>(), fna3d_64x128x32_sm75_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>(), fna3d_32x128x64_sm75_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_64x64x64_sm75_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>(), fna3d_64x128x64_sm75_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>(), fna3d_32x128x128_sm75_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_64x64x128_sm75_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>(), fna3d_64x128x128_sm75_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>(), fna3d_32x128x65536_sm75_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>(), fna3d_64x128x65536_sm75_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_64x64x65536_sm75_float16_cm_1_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm75_float16_cm_1_1_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>(), fna3d_32x128x32_sm75_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_64x64x32_sm75_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>(), fna3d_64x128x32_sm75_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>(), fna3d_32x128x64_sm75_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_64x64x64_sm75_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>(), fna3d_64x128x64_sm75_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>(), fna3d_32x128x128_sm75_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_64x64x128_sm75_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>(), fna3d_64x128x128_sm75_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>(), fna3d_32x128x65536_sm75_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>(), fna3d_64x128x65536_sm75_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_64x64x65536_sm75_float16_cm_1_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm75_float16_cm_1_1_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>(), fna3d_32x128x32_sm75_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_64x64x32_sm75_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>(), fna3d_64x128x32_sm75_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>(), fna3d_32x128x64_sm75_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_64x64x64_sm75_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>(), fna3d_64x128x64_sm75_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>(), fna3d_32x128x128_sm75_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_64x64x128_sm75_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>(), fna3d_64x128x128_sm75_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>(), fna3d_32x128x65536_sm75_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>(), fna3d_64x128x65536_sm75_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_64x64x65536_sm75_float16_cm_1_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm80_float32_cm_0_0_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>(), fna3d_32x128x32_sm80_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_64x64x32_sm80_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>(), fna3d_64x128x32_sm80_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>(), fna3d_32x128x64_sm80_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_64x64x64_sm80_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>(), fna3d_64x128x64_sm80_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>(), fna3d_32x128x128_sm80_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_64x64x128_sm80_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>(), fna3d_64x128x128_sm80_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna3d_32x128x65536_sm80_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna3d_64x128x65536_sm80_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_64x64x65536_sm80_float32_cm_0_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm80_float32_cm_0_0_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>(), fna3d_32x128x32_sm80_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_64x64x32_sm80_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>(), fna3d_64x128x32_sm80_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>(), fna3d_32x128x64_sm80_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_64x64x64_sm80_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>(), fna3d_64x128x64_sm80_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>(), fna3d_32x128x128_sm80_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_64x64x128_sm80_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>(), fna3d_64x128x128_sm80_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna3d_32x128x65536_sm80_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna3d_64x128x65536_sm80_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_64x64x65536_sm80_float32_cm_0_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm80_float32_cm_0_1_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>(), fna3d_32x128x32_sm80_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_64x64x32_sm80_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>(), fna3d_64x128x32_sm80_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>(), fna3d_32x128x64_sm80_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_64x64x64_sm80_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>(), fna3d_64x128x64_sm80_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>(), fna3d_32x128x128_sm80_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_64x64x128_sm80_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>(), fna3d_64x128x128_sm80_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna3d_32x128x65536_sm80_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna3d_64x128x65536_sm80_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_64x64x65536_sm80_float32_cm_0_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm80_float32_cm_0_1_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>(), fna3d_32x128x32_sm80_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_64x64x32_sm80_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>(), fna3d_64x128x32_sm80_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>(), fna3d_32x128x64_sm80_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_64x64x64_sm80_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>(), fna3d_64x128x64_sm80_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>(), fna3d_32x128x128_sm80_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_64x64x128_sm80_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>(), fna3d_64x128x128_sm80_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna3d_32x128x65536_sm80_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna3d_64x128x65536_sm80_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_64x64x65536_sm80_float32_cm_0_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm80_float32_cm_1_0_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>(), fna3d_32x128x32_sm80_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_64x64x32_sm80_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>(), fna3d_64x128x32_sm80_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>(), fna3d_32x128x64_sm80_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_64x64x64_sm80_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>(), fna3d_64x128x64_sm80_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>(), fna3d_32x128x128_sm80_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_64x64x128_sm80_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>(), fna3d_64x128x128_sm80_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna3d_32x128x65536_sm80_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna3d_64x128x65536_sm80_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_64x64x65536_sm80_float32_cm_1_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm80_float32_cm_1_0_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>(), fna3d_32x128x32_sm80_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_64x64x32_sm80_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>(), fna3d_64x128x32_sm80_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>(), fna3d_32x128x64_sm80_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_64x64x64_sm80_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>(), fna3d_64x128x64_sm80_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>(), fna3d_32x128x128_sm80_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_64x64x128_sm80_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>(), fna3d_64x128x128_sm80_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna3d_32x128x65536_sm80_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna3d_64x128x65536_sm80_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_64x64x65536_sm80_float32_cm_1_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm80_float32_cm_1_1_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>(), fna3d_32x128x32_sm80_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_64x64x32_sm80_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>(), fna3d_64x128x32_sm80_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>(), fna3d_32x128x64_sm80_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_64x64x64_sm80_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>(), fna3d_64x128x64_sm80_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>(), fna3d_32x128x128_sm80_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_64x64x128_sm80_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>(), fna3d_64x128x128_sm80_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna3d_32x128x65536_sm80_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna3d_64x128x65536_sm80_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_64x64x65536_sm80_float32_cm_1_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm80_float32_cm_1_1_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>(), fna3d_32x128x32_sm80_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_64x64x32_sm80_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>(), fna3d_64x128x32_sm80_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>(), fna3d_32x128x64_sm80_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_64x64x64_sm80_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>(), fna3d_64x128x64_sm80_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>(), fna3d_32x128x128_sm80_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_64x64x128_sm80_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>(), fna3d_64x128x128_sm80_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna3d_32x128x65536_sm80_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna3d_64x128x65536_sm80_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_64x64x65536_sm80_float32_cm_1_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm80_float16_cm_0_0_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>(), fna3d_32x128x32_sm80_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_64x64x32_sm80_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>(), fna3d_64x128x32_sm80_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>(), fna3d_32x128x64_sm80_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_64x64x64_sm80_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>(), fna3d_64x128x64_sm80_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>(), fna3d_32x128x128_sm80_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_64x64x128_sm80_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>(), fna3d_64x128x128_sm80_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna3d_32x128x65536_sm80_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna3d_64x128x65536_sm80_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_64x64x65536_sm80_float16_cm_0_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm80_float16_cm_0_0_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>(), fna3d_32x128x32_sm80_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_64x64x32_sm80_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>(), fna3d_64x128x32_sm80_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>(), fna3d_32x128x64_sm80_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_64x64x64_sm80_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>(), fna3d_64x128x64_sm80_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>(), fna3d_32x128x128_sm80_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_64x64x128_sm80_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>(), fna3d_64x128x128_sm80_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna3d_32x128x65536_sm80_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna3d_64x128x65536_sm80_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_64x64x65536_sm80_float16_cm_0_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm80_float16_cm_0_1_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>(), fna3d_32x128x32_sm80_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_64x64x32_sm80_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>(), fna3d_64x128x32_sm80_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>(), fna3d_32x128x64_sm80_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_64x64x64_sm80_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>(), fna3d_64x128x64_sm80_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>(), fna3d_32x128x128_sm80_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_64x64x128_sm80_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>(), fna3d_64x128x128_sm80_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna3d_32x128x65536_sm80_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna3d_64x128x65536_sm80_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_64x64x65536_sm80_float16_cm_0_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm80_float16_cm_0_1_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>(), fna3d_32x128x32_sm80_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_64x64x32_sm80_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>(), fna3d_64x128x32_sm80_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>(), fna3d_32x128x64_sm80_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_64x64x64_sm80_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>(), fna3d_64x128x64_sm80_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>(), fna3d_32x128x128_sm80_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_64x64x128_sm80_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>(), fna3d_64x128x128_sm80_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna3d_32x128x65536_sm80_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna3d_64x128x65536_sm80_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_64x64x65536_sm80_float16_cm_0_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm80_float16_cm_1_0_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>(), fna3d_32x128x32_sm80_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_64x64x32_sm80_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>(), fna3d_64x128x32_sm80_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>(), fna3d_32x128x64_sm80_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_64x64x64_sm80_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>(), fna3d_64x128x64_sm80_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>(), fna3d_32x128x128_sm80_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_64x64x128_sm80_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>(), fna3d_64x128x128_sm80_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna3d_32x128x65536_sm80_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna3d_64x128x65536_sm80_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_64x64x65536_sm80_float16_cm_1_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm80_float16_cm_1_0_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>(), fna3d_32x128x32_sm80_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_64x64x32_sm80_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>(), fna3d_64x128x32_sm80_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>(), fna3d_32x128x64_sm80_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_64x64x64_sm80_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>(), fna3d_64x128x64_sm80_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>(), fna3d_32x128x128_sm80_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_64x64x128_sm80_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>(), fna3d_64x128x128_sm80_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna3d_32x128x65536_sm80_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna3d_64x128x65536_sm80_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_64x64x65536_sm80_float16_cm_1_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm80_float16_cm_1_1_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>(), fna3d_32x128x32_sm80_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_64x64x32_sm80_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>(), fna3d_64x128x32_sm80_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>(), fna3d_32x128x64_sm80_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_64x64x64_sm80_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>(), fna3d_64x128x64_sm80_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>(), fna3d_32x128x128_sm80_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_64x64x128_sm80_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>(), fna3d_64x128x128_sm80_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna3d_32x128x65536_sm80_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna3d_64x128x65536_sm80_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_64x64x65536_sm80_float16_cm_1_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm80_float16_cm_1_1_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>(), fna3d_32x128x32_sm80_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_64x64x32_sm80_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>(), fna3d_64x128x32_sm80_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>(), fna3d_32x128x64_sm80_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_64x64x64_sm80_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>(), fna3d_64x128x64_sm80_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>(), fna3d_32x128x128_sm80_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_64x64x128_sm80_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>(), fna3d_64x128x128_sm80_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna3d_32x128x65536_sm80_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna3d_64x128x65536_sm80_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_64x64x65536_sm80_float16_cm_1_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm80_bfloat16_cm_0_0_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>(), fna3d_32x128x32_sm80_bfloat16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_64x64x32_sm80_bfloat16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>(), fna3d_64x128x32_sm80_bfloat16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>(), fna3d_32x128x64_sm80_bfloat16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_64x64x64_sm80_bfloat16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>(), fna3d_64x128x64_sm80_bfloat16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>(), fna3d_32x128x128_sm80_bfloat16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_64x64x128_sm80_bfloat16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>(), fna3d_64x128x128_sm80_bfloat16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna3d_32x128x65536_sm80_bfloat16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna3d_64x128x65536_sm80_bfloat16_cm_0_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_64x64x65536_sm80_bfloat16_cm_0_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm80_bfloat16_cm_0_0_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>(), fna3d_32x128x32_sm80_bfloat16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_64x64x32_sm80_bfloat16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>(), fna3d_64x128x32_sm80_bfloat16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>(), fna3d_32x128x64_sm80_bfloat16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_64x64x64_sm80_bfloat16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>(), fna3d_64x128x64_sm80_bfloat16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>(), fna3d_32x128x128_sm80_bfloat16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_64x64x128_sm80_bfloat16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>(), fna3d_64x128x128_sm80_bfloat16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna3d_32x128x65536_sm80_bfloat16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna3d_64x128x65536_sm80_bfloat16_cm_0_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_64x64x65536_sm80_bfloat16_cm_0_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm80_bfloat16_cm_0_1_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>(), fna3d_32x128x32_sm80_bfloat16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_64x64x32_sm80_bfloat16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>(), fna3d_64x128x32_sm80_bfloat16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>(), fna3d_32x128x64_sm80_bfloat16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_64x64x64_sm80_bfloat16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>(), fna3d_64x128x64_sm80_bfloat16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>(), fna3d_32x128x128_sm80_bfloat16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_64x64x128_sm80_bfloat16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>(), fna3d_64x128x128_sm80_bfloat16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna3d_32x128x65536_sm80_bfloat16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna3d_64x128x65536_sm80_bfloat16_cm_0_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_64x64x65536_sm80_bfloat16_cm_0_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm80_bfloat16_cm_0_1_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>(), fna3d_32x128x32_sm80_bfloat16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_64x64x32_sm80_bfloat16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>(), fna3d_64x128x32_sm80_bfloat16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>(), fna3d_32x128x64_sm80_bfloat16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_64x64x64_sm80_bfloat16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>(), fna3d_64x128x64_sm80_bfloat16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>(), fna3d_32x128x128_sm80_bfloat16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_64x64x128_sm80_bfloat16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>(), fna3d_64x128x128_sm80_bfloat16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna3d_32x128x65536_sm80_bfloat16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna3d_64x128x65536_sm80_bfloat16_cm_0_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_64x64x65536_sm80_bfloat16_cm_0_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm80_bfloat16_cm_1_0_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>(), fna3d_32x128x32_sm80_bfloat16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_64x64x32_sm80_bfloat16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>(), fna3d_64x128x32_sm80_bfloat16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>(), fna3d_32x128x64_sm80_bfloat16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_64x64x64_sm80_bfloat16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>(), fna3d_64x128x64_sm80_bfloat16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>(), fna3d_32x128x128_sm80_bfloat16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_64x64x128_sm80_bfloat16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>(), fna3d_64x128x128_sm80_bfloat16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna3d_32x128x65536_sm80_bfloat16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna3d_64x128x65536_sm80_bfloat16_cm_1_0_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_64x64x65536_sm80_bfloat16_cm_1_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm80_bfloat16_cm_1_0_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>(), fna3d_32x128x32_sm80_bfloat16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_64x64x32_sm80_bfloat16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>(), fna3d_64x128x32_sm80_bfloat16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>(), fna3d_32x128x64_sm80_bfloat16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_64x64x64_sm80_bfloat16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>(), fna3d_64x128x64_sm80_bfloat16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>(), fna3d_32x128x128_sm80_bfloat16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_64x64x128_sm80_bfloat16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>(), fna3d_64x128x128_sm80_bfloat16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna3d_32x128x65536_sm80_bfloat16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna3d_64x128x65536_sm80_bfloat16_cm_1_0_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_64x64x65536_sm80_bfloat16_cm_1_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm80_bfloat16_cm_1_1_0(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>(), fna3d_32x128x32_sm80_bfloat16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_64x64x32_sm80_bfloat16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>(), fna3d_64x128x32_sm80_bfloat16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>(), fna3d_32x128x64_sm80_bfloat16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_64x64x64_sm80_bfloat16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>(), fna3d_64x128x64_sm80_bfloat16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>(), fna3d_32x128x128_sm80_bfloat16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_64x64x128_sm80_bfloat16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>(), fna3d_64x128x128_sm80_bfloat16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna3d_32x128x65536_sm80_bfloat16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna3d_64x128x65536_sm80_bfloat16_cm_1_1_0); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_64x64x65536_sm80_bfloat16_cm_1_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +template +void fna3d_sm80_bfloat16_cm_1_1_1(T cb) { + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>(), fna3d_32x128x32_sm80_bfloat16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_64x64x32_sm80_bfloat16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>(), fna3d_64x128x32_sm80_bfloat16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>(), fna3d_32x128x64_sm80_bfloat16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_64x64x64_sm80_bfloat16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>(), fna3d_64x128x64_sm80_bfloat16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>(), fna3d_32x128x128_sm80_bfloat16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_64x64x128_sm80_bfloat16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>(), fna3d_64x128x128_sm80_bfloat16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>(), fna3d_32x128x65536_sm80_bfloat16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>(), fna3d_64x128x65536_sm80_bfloat16_cm_1_1_1); + cb(FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_64x64x65536_sm80_bfloat16_cm_1_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm50_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm50_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm50_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm50_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna1d_backward_sm50_float32_cm_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>(), fna1d_backward_64x64x32_sm50_float32_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>(), fna1d_backward_64x64x64_sm50_float32_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>(), fna1d_backward_64x64x128_sm50_float32_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna1d_backward_64x64x65536_sm50_float32_cm_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm50_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm50_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm50_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm50_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna1d_backward_sm50_float32_cm_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>(), fna1d_backward_64x64x32_sm50_float32_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>(), fna1d_backward_64x64x64_sm50_float32_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>(), fna1d_backward_64x64x128_sm50_float32_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna1d_backward_64x64x65536_sm50_float32_cm_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm50_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm50_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm50_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm50_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna1d_backward_sm50_float16_cm_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>(), fna1d_backward_64x64x32_sm50_float16_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>(), fna1d_backward_64x64x64_sm50_float16_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>(), fna1d_backward_64x64x128_sm50_float16_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna1d_backward_64x64x65536_sm50_float16_cm_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm50_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm50_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm50_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm50_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna1d_backward_sm50_float16_cm_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>(), fna1d_backward_64x64x32_sm50_float16_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>(), fna1d_backward_64x64x64_sm50_float16_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>(), fna1d_backward_64x64x128_sm50_float16_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna1d_backward_64x64x65536_sm50_float16_cm_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm70_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm70_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm70_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm70_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna1d_backward_sm70_float32_cm_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>(), fna1d_backward_64x64x32_sm70_float32_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>(), fna1d_backward_64x64x64_sm70_float32_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>(), fna1d_backward_64x64x128_sm70_float32_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna1d_backward_64x64x65536_sm70_float32_cm_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm70_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm70_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm70_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm70_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna1d_backward_sm70_float32_cm_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>(), fna1d_backward_64x64x32_sm70_float32_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>(), fna1d_backward_64x64x64_sm70_float32_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>(), fna1d_backward_64x64x128_sm70_float32_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna1d_backward_64x64x65536_sm70_float32_cm_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm70_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm70_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm70_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kMinBlocksPerSm) +fna1d_backward_128x64x128_sm70_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm70_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_128x64x65536_sm70_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::Params p); + + +template +void fna1d_backward_sm70_float16_cm_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>(), fna1d_backward_64x64x32_sm70_float16_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>(), fna1d_backward_64x64x64_sm70_float16_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>(), fna1d_backward_64x64x128_sm70_float16_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>(), fna1d_backward_128x64x128_sm70_float16_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna1d_backward_64x64x65536_sm70_float16_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>(), fna1d_backward_128x64x65536_sm70_float16_cm_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm70_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm70_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm70_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kMinBlocksPerSm) +fna1d_backward_128x64x128_sm70_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm70_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_128x64x65536_sm70_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::Params p); + + +template +void fna1d_backward_sm70_float16_cm_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>(), fna1d_backward_64x64x32_sm70_float16_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>(), fna1d_backward_64x64x64_sm70_float16_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>(), fna1d_backward_64x64x128_sm70_float16_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>(), fna1d_backward_128x64x128_sm70_float16_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna1d_backward_64x64x65536_sm70_float16_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>(), fna1d_backward_128x64x65536_sm70_float16_cm_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm75_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm75_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm75_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm75_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna1d_backward_sm75_float32_cm_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>(), fna1d_backward_64x64x32_sm75_float32_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>(), fna1d_backward_64x64x64_sm75_float32_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>(), fna1d_backward_64x64x128_sm75_float32_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna1d_backward_64x64x65536_sm75_float32_cm_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm75_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm75_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm75_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm75_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna1d_backward_sm75_float32_cm_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>(), fna1d_backward_64x64x32_sm75_float32_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>(), fna1d_backward_64x64x64_sm75_float32_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>(), fna1d_backward_64x64x128_sm75_float32_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna1d_backward_64x64x65536_sm75_float32_cm_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm75_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm75_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm75_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm75_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna1d_backward_sm75_float16_cm_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>(), fna1d_backward_64x64x32_sm75_float16_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>(), fna1d_backward_64x64x64_sm75_float16_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>(), fna1d_backward_64x64x128_sm75_float16_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna1d_backward_64x64x65536_sm75_float16_cm_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm75_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm75_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm75_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm75_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna1d_backward_sm75_float16_cm_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>(), fna1d_backward_64x64x32_sm75_float16_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>(), fna1d_backward_64x64x64_sm75_float16_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>(), fna1d_backward_64x64x128_sm75_float16_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna1d_backward_64x64x65536_sm75_float16_cm_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm80_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm80_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm80_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kMinBlocksPerSm) +fna1d_backward_128x64x128_sm80_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm80_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_128x64x65536_sm80_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna1d_backward_sm80_float32_cm_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>(), fna1d_backward_64x64x32_sm80_float32_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>(), fna1d_backward_64x64x64_sm80_float32_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>(), fna1d_backward_64x64x128_sm80_float32_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>(), fna1d_backward_128x64x128_sm80_float32_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna1d_backward_64x64x65536_sm80_float32_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna1d_backward_128x64x65536_sm80_float32_cm_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm80_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm80_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm80_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kMinBlocksPerSm) +fna1d_backward_128x64x128_sm80_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm80_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_128x64x65536_sm80_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna1d_backward_sm80_float32_cm_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>(), fna1d_backward_64x64x32_sm80_float32_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>(), fna1d_backward_64x64x64_sm80_float32_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>(), fna1d_backward_64x64x128_sm80_float32_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>(), fna1d_backward_128x64x128_sm80_float32_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna1d_backward_64x64x65536_sm80_float32_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna1d_backward_128x64x65536_sm80_float32_cm_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm80_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm80_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm80_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna1d_backward_128x128x128_sm80_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm80_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_128x64x65536_sm80_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna1d_backward_sm80_float16_cm_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna1d_backward_64x64x32_sm80_float16_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna1d_backward_64x64x64_sm80_float16_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna1d_backward_64x64x128_sm80_float16_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>(), fna1d_backward_128x128x128_sm80_float16_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna1d_backward_64x64x65536_sm80_float16_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna1d_backward_128x64x65536_sm80_float16_cm_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm80_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm80_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm80_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna1d_backward_128x128x128_sm80_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm80_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_128x64x65536_sm80_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna1d_backward_sm80_float16_cm_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna1d_backward_64x64x32_sm80_float16_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna1d_backward_64x64x64_sm80_float16_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna1d_backward_64x64x128_sm80_float16_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>(), fna1d_backward_128x128x128_sm80_float16_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna1d_backward_64x64x65536_sm80_float16_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna1d_backward_128x64x65536_sm80_float16_cm_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / bfloat16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna1d_backward_128x128x128_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_128x64x65536_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna1d_backward_sm80_bfloat16_cm_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna1d_backward_64x64x32_sm80_bfloat16_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna1d_backward_64x64x64_sm80_bfloat16_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna1d_backward_64x64x128_sm80_bfloat16_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>(), fna1d_backward_128x128x128_sm80_bfloat16_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna1d_backward_64x64x65536_sm80_bfloat16_cm_0); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna1d_backward_128x64x65536_sm80_bfloat16_cm_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-1D / bfloat16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna1d_backward_128x128x128_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_128x64x65536_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna1d_backward_sm80_bfloat16_cm_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna1d_backward_64x64x32_sm80_bfloat16_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna1d_backward_64x64x64_sm80_bfloat16_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna1d_backward_64x64x128_sm80_bfloat16_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>(), fna1d_backward_128x128x128_sm80_bfloat16_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna1d_backward_64x64x65536_sm80_bfloat16_cm_1); + cb(FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna1d_backward_128x64x65536_sm80_bfloat16_cm_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna2d_backward_sm50_float32_cm_0_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm50_float32_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm50_float32_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm50_float32_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm50_float32_cm_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna2d_backward_sm50_float32_cm_0_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm50_float32_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm50_float32_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm50_float32_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm50_float32_cm_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna2d_backward_sm50_float32_cm_1_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm50_float32_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm50_float32_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm50_float32_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm50_float32_cm_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna2d_backward_sm50_float32_cm_1_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm50_float32_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm50_float32_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm50_float32_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm50_float32_cm_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna2d_backward_sm50_float16_cm_0_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm50_float16_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm50_float16_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm50_float16_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm50_float16_cm_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna2d_backward_sm50_float16_cm_0_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm50_float16_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm50_float16_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm50_float16_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm50_float16_cm_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna2d_backward_sm50_float16_cm_1_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm50_float16_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm50_float16_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm50_float16_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm50_float16_cm_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna2d_backward_sm50_float16_cm_1_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm50_float16_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm50_float16_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm50_float16_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm50_float16_cm_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna2d_backward_sm70_float32_cm_0_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm70_float32_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm70_float32_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm70_float32_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm70_float32_cm_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna2d_backward_sm70_float32_cm_0_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm70_float32_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm70_float32_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm70_float32_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm70_float32_cm_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna2d_backward_sm70_float32_cm_1_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm70_float32_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm70_float32_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm70_float32_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm70_float32_cm_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna2d_backward_sm70_float32_cm_1_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm70_float32_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm70_float32_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm70_float32_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm70_float32_cm_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kMinBlocksPerSm) +fna2d_backward_128x64x128_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::Params p); + + +template +void fna2d_backward_sm70_float16_cm_0_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm70_float16_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm70_float16_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm70_float16_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>(), fna2d_backward_128x64x128_sm70_float16_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm70_float16_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>(), fna2d_backward_128x64x65536_sm70_float16_cm_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kMinBlocksPerSm) +fna2d_backward_128x64x128_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::Params p); + + +template +void fna2d_backward_sm70_float16_cm_0_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm70_float16_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm70_float16_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm70_float16_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>(), fna2d_backward_128x64x128_sm70_float16_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm70_float16_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>(), fna2d_backward_128x64x65536_sm70_float16_cm_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kMinBlocksPerSm) +fna2d_backward_128x64x128_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::Params p); + + +template +void fna2d_backward_sm70_float16_cm_1_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm70_float16_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm70_float16_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm70_float16_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>(), fna2d_backward_128x64x128_sm70_float16_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm70_float16_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>(), fna2d_backward_128x64x65536_sm70_float16_cm_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kMinBlocksPerSm) +fna2d_backward_128x64x128_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::Params p); + + +template +void fna2d_backward_sm70_float16_cm_1_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm70_float16_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm70_float16_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm70_float16_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>(), fna2d_backward_128x64x128_sm70_float16_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm70_float16_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>(), fna2d_backward_128x64x65536_sm70_float16_cm_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna2d_backward_sm75_float32_cm_0_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm75_float32_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm75_float32_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm75_float32_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm75_float32_cm_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna2d_backward_sm75_float32_cm_0_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm75_float32_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm75_float32_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm75_float32_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm75_float32_cm_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna2d_backward_sm75_float32_cm_1_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm75_float32_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm75_float32_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm75_float32_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm75_float32_cm_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna2d_backward_sm75_float32_cm_1_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm75_float32_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm75_float32_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm75_float32_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm75_float32_cm_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna2d_backward_sm75_float16_cm_0_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm75_float16_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm75_float16_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm75_float16_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm75_float16_cm_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna2d_backward_sm75_float16_cm_0_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm75_float16_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm75_float16_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm75_float16_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm75_float16_cm_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna2d_backward_sm75_float16_cm_1_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm75_float16_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm75_float16_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm75_float16_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm75_float16_cm_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna2d_backward_sm75_float16_cm_1_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm75_float16_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm75_float16_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm75_float16_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm75_float16_cm_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kMinBlocksPerSm) +fna2d_backward_128x64x128_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna2d_backward_sm80_float32_cm_0_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm80_float32_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm80_float32_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm80_float32_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>(), fna2d_backward_128x64x128_sm80_float32_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm80_float32_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna2d_backward_128x64x65536_sm80_float32_cm_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kMinBlocksPerSm) +fna2d_backward_128x64x128_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna2d_backward_sm80_float32_cm_0_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm80_float32_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm80_float32_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm80_float32_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>(), fna2d_backward_128x64x128_sm80_float32_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm80_float32_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna2d_backward_128x64x65536_sm80_float32_cm_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kMinBlocksPerSm) +fna2d_backward_128x64x128_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna2d_backward_sm80_float32_cm_1_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm80_float32_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm80_float32_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm80_float32_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>(), fna2d_backward_128x64x128_sm80_float32_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm80_float32_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna2d_backward_128x64x65536_sm80_float32_cm_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kMinBlocksPerSm) +fna2d_backward_128x64x128_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna2d_backward_sm80_float32_cm_1_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm80_float32_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm80_float32_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm80_float32_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>(), fna2d_backward_128x64x128_sm80_float32_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm80_float32_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna2d_backward_128x64x65536_sm80_float32_cm_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna2d_backward_128x128x128_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna2d_backward_sm80_float16_cm_0_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm80_float16_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm80_float16_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm80_float16_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>(), fna2d_backward_128x128x128_sm80_float16_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm80_float16_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna2d_backward_128x64x65536_sm80_float16_cm_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna2d_backward_128x128x128_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna2d_backward_sm80_float16_cm_0_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm80_float16_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm80_float16_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm80_float16_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>(), fna2d_backward_128x128x128_sm80_float16_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm80_float16_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna2d_backward_128x64x65536_sm80_float16_cm_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna2d_backward_128x128x128_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna2d_backward_sm80_float16_cm_1_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm80_float16_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm80_float16_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm80_float16_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>(), fna2d_backward_128x128x128_sm80_float16_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm80_float16_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna2d_backward_128x64x65536_sm80_float16_cm_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna2d_backward_128x128x128_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna2d_backward_sm80_float16_cm_1_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm80_float16_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm80_float16_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm80_float16_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>(), fna2d_backward_128x128x128_sm80_float16_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm80_float16_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna2d_backward_128x64x65536_sm80_float16_cm_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / bfloat16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna2d_backward_128x128x128_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna2d_backward_sm80_bfloat16_cm_0_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm80_bfloat16_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm80_bfloat16_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm80_bfloat16_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>(), fna2d_backward_128x128x128_sm80_bfloat16_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm80_bfloat16_cm_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna2d_backward_128x64x65536_sm80_bfloat16_cm_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / bfloat16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna2d_backward_128x128x128_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna2d_backward_sm80_bfloat16_cm_0_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm80_bfloat16_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm80_bfloat16_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm80_bfloat16_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>(), fna2d_backward_128x128x128_sm80_bfloat16_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm80_bfloat16_cm_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna2d_backward_128x64x65536_sm80_bfloat16_cm_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / bfloat16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna2d_backward_128x128x128_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna2d_backward_sm80_bfloat16_cm_1_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm80_bfloat16_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm80_bfloat16_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm80_bfloat16_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>(), fna2d_backward_128x128x128_sm80_bfloat16_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm80_bfloat16_cm_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna2d_backward_128x64x65536_sm80_bfloat16_cm_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-2D / bfloat16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna2d_backward_128x128x128_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna2d_backward_sm80_bfloat16_cm_1_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna2d_backward_64x64x32_sm80_bfloat16_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna2d_backward_64x64x64_sm80_bfloat16_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna2d_backward_64x64x128_sm80_bfloat16_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>(), fna2d_backward_128x128x128_sm80_bfloat16_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna2d_backward_64x64x65536_sm80_bfloat16_cm_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna2d_backward_128x64x65536_sm80_bfloat16_cm_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm50_float32_cm_0_0_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm50_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm50_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm50_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm50_float32_cm_0_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm50_float32_cm_0_0_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm50_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm50_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm50_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm50_float32_cm_0_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm50_float32_cm_0_1_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm50_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm50_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm50_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm50_float32_cm_0_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm50_float32_cm_0_1_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm50_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm50_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm50_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm50_float32_cm_0_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm50_float32_cm_1_0_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm50_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm50_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm50_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm50_float32_cm_1_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm50_float32_cm_1_0_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm50_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm50_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm50_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm50_float32_cm_1_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm50_float32_cm_1_1_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm50_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm50_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm50_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm50_float32_cm_1_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm50_float32_cm_1_1_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm50_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm50_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm50_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm50_float32_cm_1_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm50_float16_cm_0_0_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm50_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm50_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm50_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm50_float16_cm_0_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm50_float16_cm_0_0_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm50_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm50_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm50_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm50_float16_cm_0_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm50_float16_cm_0_1_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm50_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm50_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm50_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm50_float16_cm_0_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm50_float16_cm_0_1_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm50_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm50_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm50_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm50_float16_cm_0_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm50_float16_cm_1_0_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm50_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm50_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm50_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm50_float16_cm_1_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm50_float16_cm_1_0_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm50_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm50_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm50_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm50_float16_cm_1_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm50_float16_cm_1_1_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm50_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm50_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm50_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm50_float16_cm_1_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm50_float16_cm_1_1_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm50_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm50_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm50_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm50_float16_cm_1_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm70_float32_cm_0_0_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm70_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm70_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm70_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm70_float32_cm_0_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm70_float32_cm_0_0_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm70_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm70_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm70_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm70_float32_cm_0_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm70_float32_cm_0_1_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm70_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm70_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm70_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm70_float32_cm_0_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm70_float32_cm_0_1_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm70_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm70_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm70_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm70_float32_cm_0_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm70_float32_cm_1_0_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm70_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm70_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm70_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm70_float32_cm_1_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm70_float32_cm_1_0_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm70_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm70_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm70_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm70_float32_cm_1_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm70_float32_cm_1_1_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm70_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm70_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm70_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm70_float32_cm_1_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm70_float32_cm_1_1_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm70_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm70_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm70_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm70_float32_cm_1_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm70_float16_cm_0_0_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm70_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm70_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm70_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>(), fna3d_backward_128x64x128_sm70_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm70_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm70_float16_cm_0_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm70_float16_cm_0_0_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm70_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm70_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm70_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>(), fna3d_backward_128x64x128_sm70_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm70_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm70_float16_cm_0_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm70_float16_cm_0_1_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm70_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm70_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm70_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>(), fna3d_backward_128x64x128_sm70_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm70_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm70_float16_cm_0_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm70_float16_cm_0_1_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm70_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm70_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm70_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>(), fna3d_backward_128x64x128_sm70_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm70_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm70_float16_cm_0_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm70_float16_cm_1_0_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm70_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm70_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm70_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>(), fna3d_backward_128x64x128_sm70_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm70_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm70_float16_cm_1_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm70_float16_cm_1_0_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm70_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm70_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm70_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>(), fna3d_backward_128x64x128_sm70_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm70_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm70_float16_cm_1_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm70_float16_cm_1_1_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm70_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm70_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm70_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>(), fna3d_backward_128x64x128_sm70_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm70_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm70_float16_cm_1_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm70_float16_cm_1_1_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm70_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm70_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm70_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>(), fna3d_backward_128x64x128_sm70_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm70_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm70_float16_cm_1_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm75_float32_cm_0_0_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm75_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm75_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm75_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm75_float32_cm_0_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm75_float32_cm_0_0_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm75_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm75_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm75_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm75_float32_cm_0_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm75_float32_cm_0_1_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm75_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm75_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm75_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm75_float32_cm_0_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm75_float32_cm_0_1_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm75_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm75_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm75_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm75_float32_cm_0_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm75_float32_cm_1_0_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm75_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm75_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm75_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm75_float32_cm_1_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm75_float32_cm_1_0_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm75_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm75_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm75_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm75_float32_cm_1_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm75_float32_cm_1_1_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm75_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm75_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm75_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm75_float32_cm_1_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm75_float32_cm_1_1_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm75_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm75_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm75_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm75_float32_cm_1_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm75_float16_cm_0_0_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm75_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm75_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm75_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm75_float16_cm_0_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm75_float16_cm_0_0_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm75_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm75_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm75_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm75_float16_cm_0_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm75_float16_cm_0_1_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm75_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm75_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm75_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm75_float16_cm_0_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm75_float16_cm_0_1_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm75_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm75_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm75_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm75_float16_cm_0_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm75_float16_cm_1_0_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm75_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm75_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm75_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm75_float16_cm_1_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm75_float16_cm_1_0_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm75_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm75_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm75_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm75_float16_cm_1_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm75_float16_cm_1_1_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm75_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm75_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm75_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm75_float16_cm_1_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p); + + +template +void fna3d_backward_sm75_float16_cm_1_1_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm75_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm75_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm75_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm75_float16_cm_1_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm80_float32_cm_0_0_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm80_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm80_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm80_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>(), fna3d_backward_128x64x128_sm80_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm80_float32_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm80_float32_cm_0_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm80_float32_cm_0_0_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm80_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm80_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm80_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>(), fna3d_backward_128x64x128_sm80_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm80_float32_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm80_float32_cm_0_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm80_float32_cm_0_1_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm80_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm80_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm80_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>(), fna3d_backward_128x64x128_sm80_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm80_float32_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm80_float32_cm_0_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm80_float32_cm_0_1_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm80_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm80_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm80_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>(), fna3d_backward_128x64x128_sm80_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm80_float32_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm80_float32_cm_0_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm80_float32_cm_1_0_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm80_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm80_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm80_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>(), fna3d_backward_128x64x128_sm80_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm80_float32_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm80_float32_cm_1_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm80_float32_cm_1_0_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm80_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm80_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm80_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>(), fna3d_backward_128x64x128_sm80_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm80_float32_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm80_float32_cm_1_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm80_float32_cm_1_1_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm80_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm80_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm80_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>(), fna3d_backward_128x64x128_sm80_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm80_float32_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm80_float32_cm_1_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm80_float32_cm_1_1_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm80_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm80_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm80_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>(), fna3d_backward_128x64x128_sm80_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm80_float32_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm80_float32_cm_1_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm80_float16_cm_0_0_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm80_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm80_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm80_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>(), fna3d_backward_128x128x128_sm80_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm80_float16_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm80_float16_cm_0_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm80_float16_cm_0_0_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm80_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm80_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm80_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>(), fna3d_backward_128x128x128_sm80_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm80_float16_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm80_float16_cm_0_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm80_float16_cm_0_1_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm80_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm80_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm80_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>(), fna3d_backward_128x128x128_sm80_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm80_float16_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm80_float16_cm_0_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm80_float16_cm_0_1_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm80_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm80_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm80_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>(), fna3d_backward_128x128x128_sm80_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm80_float16_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm80_float16_cm_0_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm80_float16_cm_1_0_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm80_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm80_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm80_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>(), fna3d_backward_128x128x128_sm80_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm80_float16_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm80_float16_cm_1_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm80_float16_cm_1_0_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm80_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm80_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm80_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>(), fna3d_backward_128x128x128_sm80_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm80_float16_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm80_float16_cm_1_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm80_float16_cm_1_1_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm80_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm80_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm80_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>(), fna3d_backward_128x128x128_sm80_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm80_float16_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm80_float16_cm_1_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm80_float16_cm_1_1_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm80_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm80_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm80_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>(), fna3d_backward_128x128x128_sm80_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm80_float16_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm80_float16_cm_1_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm80_bfloat16_cm_0_0_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm80_bfloat16_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm80_bfloat16_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm80_bfloat16_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>(), fna3d_backward_128x128x128_sm80_bfloat16_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm80_bfloat16_cm_0_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm80_bfloat16_cm_0_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm80_bfloat16_cm_0_0_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm80_bfloat16_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm80_bfloat16_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm80_bfloat16_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>(), fna3d_backward_128x128x128_sm80_bfloat16_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm80_bfloat16_cm_0_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm80_bfloat16_cm_0_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm80_bfloat16_cm_0_1_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm80_bfloat16_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm80_bfloat16_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm80_bfloat16_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>(), fna3d_backward_128x128x128_sm80_bfloat16_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm80_bfloat16_cm_0_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm80_bfloat16_cm_0_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm80_bfloat16_cm_0_1_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm80_bfloat16_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm80_bfloat16_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm80_bfloat16_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>(), fna3d_backward_128x128x128_sm80_bfloat16_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm80_bfloat16_cm_0_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm80_bfloat16_cm_0_1_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm80_bfloat16_cm_1_0_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm80_bfloat16_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm80_bfloat16_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm80_bfloat16_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>(), fna3d_backward_128x128x128_sm80_bfloat16_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm80_bfloat16_cm_1_0_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm80_bfloat16_cm_1_0_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm80_bfloat16_cm_1_0_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm80_bfloat16_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm80_bfloat16_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm80_bfloat16_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>(), fna3d_backward_128x128x128_sm80_bfloat16_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm80_bfloat16_cm_1_0_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm80_bfloat16_cm_1_0_1); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm80_bfloat16_cm_1_1_0(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm80_bfloat16_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm80_bfloat16_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm80_bfloat16_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>(), fna3d_backward_128x128x128_sm80_bfloat16_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm80_bfloat16_cm_1_1_0); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm80_bfloat16_cm_1_1_0); +} + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80Backward Kernel +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p); + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p); + + +template +void fna3d_backward_sm80_bfloat16_cm_1_1_1(T cb) { + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>(), fna3d_backward_64x64x32_sm80_bfloat16_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>(), fna3d_backward_64x64x64_sm80_bfloat16_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>(), fna3d_backward_64x64x128_sm80_bfloat16_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>(), fna3d_backward_128x128x128_sm80_bfloat16_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>(), fna3d_backward_64x64x65536_sm80_bfloat16_cm_1_1_1); + cb(FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>(), fna3d_backward_128x64x65536_sm80_bfloat16_cm_1_1_1); +} + +} // namespace natten +} // namespace cuda +} // namespace fna + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha/dispatch_head_dim.h b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha/dispatch_head_dim.h new file mode 100644 index 00000000..8fff8afc --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha/dispatch_head_dim.h @@ -0,0 +1,68 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_hopper { +#define DISPATCH_HOPPER_FMHA_FORWARD_float16(dim, q_tile_size, kv_tile_size, kernel_type, ...) \ + [&] { \ + if (dim == 32) { \ + DISPATCH_HOPPER_FMHA_FORWARD_float16_headdim32(q_tile_size, kv_tile_size, kernel_type, __VA_ARGS__); \ + } \ + else if (dim == 64) { \ + DISPATCH_HOPPER_FMHA_FORWARD_float16_headdim64(q_tile_size, kv_tile_size, kernel_type, __VA_ARGS__); \ + } \ + else if (dim == 128) { \ + DISPATCH_HOPPER_FMHA_FORWARD_float16_headdim128(q_tile_size, kv_tile_size, kernel_type, __VA_ARGS__); \ + } \ + else if (dim == 256) { \ + DISPATCH_HOPPER_FMHA_FORWARD_float16_headdim256(q_tile_size, kv_tile_size, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FMHA forward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for float16."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FMHA_FORWARD_bfloat16(dim, q_tile_size, kv_tile_size, kernel_type, ...) \ + [&] { \ + if (dim == 32) { \ + DISPATCH_HOPPER_FMHA_FORWARD_bfloat16_headdim32(q_tile_size, kv_tile_size, kernel_type, __VA_ARGS__); \ + } \ + else if (dim == 64) { \ + DISPATCH_HOPPER_FMHA_FORWARD_bfloat16_headdim64(q_tile_size, kv_tile_size, kernel_type, __VA_ARGS__); \ + } \ + else if (dim == 128) { \ + DISPATCH_HOPPER_FMHA_FORWARD_bfloat16_headdim128(q_tile_size, kv_tile_size, kernel_type, __VA_ARGS__); \ + } \ + else if (dim == 256) { \ + DISPATCH_HOPPER_FMHA_FORWARD_bfloat16_headdim256(q_tile_size, kv_tile_size, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FMHA forward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for bfloat16."); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fmha_hopper +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha/dispatch_tile_size.h b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha/dispatch_tile_size.h new file mode 100644 index 00000000..2c8af50b --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha/dispatch_tile_size.h @@ -0,0 +1,136 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_hopper { +#define DISPATCH_HOPPER_FMHA_FORWARD_float16_headdim32(q_tile_size, kv_tile_size, kernel_type, ...) \ + [&] { \ + if (q_tile_size == 64 && \ +kv_tile_size == 128 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fmha_hopper::hopper_fmha_float16_64x128x32(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FMHA forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 32): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + ", schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FMHA_FORWARD_float16_headdim64(q_tile_size, kv_tile_size, kernel_type, ...) \ + [&] { \ + if (q_tile_size == 64 && \ +kv_tile_size == 128 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fmha_hopper::hopper_fmha_float16_64x128x64(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FMHA forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 64): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + ", schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FMHA_FORWARD_float16_headdim128(q_tile_size, kv_tile_size, kernel_type, ...) \ + [&] { \ + if (q_tile_size == 128 && \ +kv_tile_size == 128 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fmha_hopper::hopper_fmha_float16_128x128x128_coop(__VA_ARGS__); \ +} \ + else if (q_tile_size == 128 && \ +kv_tile_size == 128 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fmha_hopper::hopper_fmha_float16_128x128x128_pp(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FMHA forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 128): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + ", schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FMHA_FORWARD_float16_headdim256(q_tile_size, kv_tile_size, kernel_type, ...) \ + [&] { \ + if (q_tile_size == 128 && \ +kv_tile_size == 64 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fmha_hopper::hopper_fmha_float16_128x64x256_coop(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FMHA forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 256): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + ", schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FMHA_FORWARD_bfloat16_headdim32(q_tile_size, kv_tile_size, kernel_type, ...) \ + [&] { \ + if (q_tile_size == 64 && \ +kv_tile_size == 128 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fmha_hopper::hopper_fmha_bfloat16_64x128x32(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FMHA forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 32): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + ", schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FMHA_FORWARD_bfloat16_headdim64(q_tile_size, kv_tile_size, kernel_type, ...) \ + [&] { \ + if (q_tile_size == 64 && \ +kv_tile_size == 128 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fmha_hopper::hopper_fmha_bfloat16_64x128x64(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FMHA forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 64): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + ", schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FMHA_FORWARD_bfloat16_headdim128(q_tile_size, kv_tile_size, kernel_type, ...) \ + [&] { \ + if (q_tile_size == 128 && \ +kv_tile_size == 128 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fmha_hopper::hopper_fmha_bfloat16_128x128x128_coop(__VA_ARGS__); \ +} \ + else if (q_tile_size == 128 && \ +kv_tile_size == 128 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fmha_hopper::hopper_fmha_bfloat16_128x128x128_pp(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FMHA forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 128): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + ", schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FMHA_FORWARD_bfloat16_headdim256(q_tile_size, kv_tile_size, kernel_type, ...) \ + [&] { \ + if (q_tile_size == 128 && \ +kv_tile_size == 64 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fmha_hopper::hopper_fmha_bfloat16_128x64x256_coop(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FMHA forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 256): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + ", schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fmha_hopper +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha/interface.h b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha/interface.h new file mode 100644 index 00000000..426e3e33 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha/interface.h @@ -0,0 +1,43 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_hopper { +#define DISPATCH_HOPPER_FMHA_FORWARD(dtype, dim, q_tile_size, kv_tile_size, kernel_type, ...) \ + [&] { \ + if (dtype == torch::kFloat16) { \ + DISPATCH_HOPPER_FMHA_FORWARD_float16(dim, q_tile_size, kv_tile_size, kernel_type, __VA_ARGS__); \ + } \ + else if (dtype == torch::kBFloat16) { \ + DISPATCH_HOPPER_FMHA_FORWARD_bfloat16(dim, q_tile_size, kv_tile_size, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FMHA forward kernel dispatch failed! It does not support dtype " + std::string(c10::toString(dtype)) + "."); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fmha_hopper +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha/kernels.h b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha/kernels.h new file mode 100644 index 00000000..80d345a9 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha/kernels.h @@ -0,0 +1,269 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_hopper { + +void hopper_fmha_float16_64x128x32( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fmha_float16_64x128x64( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fmha_float16_128x128x128_coop( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fmha_float16_128x128x128_pp( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fmha_float16_128x64x256_coop( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fmha_bfloat16_64x128x32( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fmha_bfloat16_64x128x64( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fmha_bfloat16_128x128x128_coop( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fmha_bfloat16_128x128x128_pp( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fmha_bfloat16_128x64x256_coop( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + + +} // namespace natten +} // namespace cuda +} // namespace fmha_hopper +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/dispatch_head_dim.h b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/dispatch_head_dim.h new file mode 100644 index 00000000..8fe7b52f --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/dispatch_head_dim.h @@ -0,0 +1,61 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_hopper { +#define DISPATCH_HOPPER_FMHA_BACKWARD_float16(dim, q_tile_size, kv_tile_size, ...) \ + [&] { \ + if (dim == 32) { \ + DISPATCH_HOPPER_FMHA_BACKWARD_float16_headdim32(q_tile_size, kv_tile_size, __VA_ARGS__); \ + } \ + else if (dim == 64) { \ + DISPATCH_HOPPER_FMHA_BACKWARD_float16_headdim64(q_tile_size, kv_tile_size, __VA_ARGS__); \ + } \ + else if (dim == 128) { \ + DISPATCH_HOPPER_FMHA_BACKWARD_float16_headdim128(q_tile_size, kv_tile_size, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FMHA backward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for float16."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FMHA_BACKWARD_bfloat16(dim, q_tile_size, kv_tile_size, ...) \ + [&] { \ + if (dim == 32) { \ + DISPATCH_HOPPER_FMHA_BACKWARD_bfloat16_headdim32(q_tile_size, kv_tile_size, __VA_ARGS__); \ + } \ + else if (dim == 64) { \ + DISPATCH_HOPPER_FMHA_BACKWARD_bfloat16_headdim64(q_tile_size, kv_tile_size, __VA_ARGS__); \ + } \ + else if (dim == 128) { \ + DISPATCH_HOPPER_FMHA_BACKWARD_bfloat16_headdim128(q_tile_size, kv_tile_size, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FMHA backward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for bfloat16."); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fmha_hopper +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/dispatch_tile_size.h b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/dispatch_tile_size.h new file mode 100644 index 00000000..bfbdb76a --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/dispatch_tile_size.h @@ -0,0 +1,111 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_hopper { +#define DISPATCH_HOPPER_FMHA_BACKWARD_float16_headdim32(q_tile_size, kv_tile_size, ...) \ + [&] { \ + if (q_tile_size == 64 && \ +kv_tile_size == 128) { \ + natten::cuda::fmha_hopper::hopper_fmha_backward_float16_64x128x32(__VA_ARGS__); \ +} \ + else if (q_tile_size == 128 && \ +kv_tile_size == 128) { \ + natten::cuda::fmha_hopper::hopper_fmha_backward_float16_128x128x32(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FMHA backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FMHA_BACKWARD_float16_headdim64(q_tile_size, kv_tile_size, ...) \ + [&] { \ + if (q_tile_size == 64 && \ +kv_tile_size == 128) { \ + natten::cuda::fmha_hopper::hopper_fmha_backward_float16_64x128x64(__VA_ARGS__); \ +} \ + else if (q_tile_size == 128 && \ +kv_tile_size == 128) { \ + natten::cuda::fmha_hopper::hopper_fmha_backward_float16_128x128x64(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FMHA backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FMHA_BACKWARD_float16_headdim128(q_tile_size, kv_tile_size, ...) \ + [&] { \ + if (q_tile_size == 64 && \ +kv_tile_size == 128) { \ + natten::cuda::fmha_hopper::hopper_fmha_backward_float16_64x128x128(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FMHA backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FMHA_BACKWARD_bfloat16_headdim32(q_tile_size, kv_tile_size, ...) \ + [&] { \ + if (q_tile_size == 64 && \ +kv_tile_size == 128) { \ + natten::cuda::fmha_hopper::hopper_fmha_backward_bfloat16_64x128x32(__VA_ARGS__); \ +} \ + else if (q_tile_size == 128 && \ +kv_tile_size == 128) { \ + natten::cuda::fmha_hopper::hopper_fmha_backward_bfloat16_128x128x32(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FMHA backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FMHA_BACKWARD_bfloat16_headdim64(q_tile_size, kv_tile_size, ...) \ + [&] { \ + if (q_tile_size == 64 && \ +kv_tile_size == 128) { \ + natten::cuda::fmha_hopper::hopper_fmha_backward_bfloat16_64x128x64(__VA_ARGS__); \ +} \ + else if (q_tile_size == 128 && \ +kv_tile_size == 128) { \ + natten::cuda::fmha_hopper::hopper_fmha_backward_bfloat16_128x128x64(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FMHA backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FMHA_BACKWARD_bfloat16_headdim128(q_tile_size, kv_tile_size, ...) \ + [&] { \ + if (q_tile_size == 64 && \ +kv_tile_size == 128) { \ + natten::cuda::fmha_hopper::hopper_fmha_backward_bfloat16_64x128x128(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FMHA backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=" + std::to_string(q_tile_size) + ", kv_tile=" + std::to_string(kv_tile_size) + "."); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fmha_hopper +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/interface.h b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/interface.h new file mode 100644 index 00000000..a168a87f --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/interface.h @@ -0,0 +1,42 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_hopper { +#define DISPATCH_HOPPER_FMHA_BACKWARD(dtype, dim, q_tile_size, kv_tile_size, ...) \ + [&] { \ + if (dtype == torch::kFloat16) { \ + DISPATCH_HOPPER_FMHA_BACKWARD_float16(dim, q_tile_size, kv_tile_size, __VA_ARGS__); \ + } \ + else if (dtype == torch::kBFloat16) { \ + DISPATCH_HOPPER_FMHA_BACKWARD_bfloat16(dim, q_tile_size, kv_tile_size, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FMHA backward kernel dispatch failed! It does not support dtype " + std::string(c10::toString(dtype)) + "."); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fmha_hopper +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/kernels.h b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/kernels.h new file mode 100644 index 00000000..e988b630 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/kernels.h @@ -0,0 +1,308 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_hopper { + +void hopper_fmha_backward_float16_64x128x32( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fmha_backward_float16_128x128x32( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fmha_backward_float16_64x128x64( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fmha_backward_float16_128x128x64( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fmha_backward_float16_64x128x128( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fmha_backward_bfloat16_64x128x32( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fmha_backward_bfloat16_128x128x32( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fmha_backward_bfloat16_64x128x64( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fmha_backward_bfloat16_128x128x64( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fmha_backward_bfloat16_64x128x128( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); + + +} // namespace natten +} // namespace cuda +} // namespace fmha_hopper +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_cm.h b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_cm.h new file mode 100644 index 00000000..291a1f27 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_cm.h @@ -0,0 +1,534 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_hopper { +#define DISPATCH_HOPPER_FNA_FORWARD_1D_float16_headdim32(is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_1D_float16_headdim32_causal0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_1D_float16_headdim32_causal1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-1D forward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_1D_float16_headdim64(is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_1D_float16_headdim64_causal0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_1D_float16_headdim64_causal1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-1D forward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_1D_float16_headdim128(is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_1D_float16_headdim128_causal0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_1D_float16_headdim128_causal1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-1D forward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_1D_float16_headdim256(is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_1D_float16_headdim256_causal0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_1D_float16_headdim256_causal1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-1D forward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 256) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_1D_bfloat16_headdim32(is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_1D_bfloat16_headdim32_causal0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_1D_bfloat16_headdim32_causal1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-1D forward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_1D_bfloat16_headdim64(is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_1D_bfloat16_headdim64_causal0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_1D_bfloat16_headdim64_causal1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-1D forward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_1D_bfloat16_headdim128(is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_1D_bfloat16_headdim128_causal0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_1D_bfloat16_headdim128_causal1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-1D forward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_1D_bfloat16_headdim256(is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_1D_bfloat16_headdim256_causal0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_1D_bfloat16_headdim256_causal1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-1D forward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 256) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim32(is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim32_causal0x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim32_causal0x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim32_causal1x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim32_causal1x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim64(is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim64_causal0x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim64_causal0x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim64_causal1x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim64_causal1x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim128(is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim128_causal0x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim128_causal0x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim128_causal1x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim128_causal1x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim256(is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim256_causal0x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim256_causal0x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim256_causal1x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim256_causal1x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 256) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim32(is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim32_causal0x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim32_causal0x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim32_causal1x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim32_causal1x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim64(is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim64_causal0x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim64_causal0x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim64_causal1x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim64_causal1x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim128(is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim128_causal0x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim128_causal0x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim128_causal1x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim128_causal1x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim256(is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim256_causal0x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim256_causal0x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim256_causal1x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim256_causal1x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 256) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim32(is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim32_causal0x0x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim32_causal0x0x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim32_causal0x1x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim32_causal0x1x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim32_causal1x0x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim32_causal1x0x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim32_causal1x1x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim32_causal1x1x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim64(is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim64_causal0x0x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim64_causal0x0x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim64_causal0x1x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim64_causal0x1x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim64_causal1x0x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim64_causal1x0x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim64_causal1x1x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim64_causal1x1x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim128(is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim128_causal0x0x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim128_causal0x0x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim128_causal0x1x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim128_causal0x1x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim128_causal1x0x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim128_causal1x0x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim128_causal1x1x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim128_causal1x1x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim256(is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim256_causal0x0x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim256_causal0x0x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim256_causal0x1x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim256_causal0x1x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim256_causal1x0x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim256_causal1x0x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim256_causal1x1x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim256_causal1x1x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 256) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim32(is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim32_causal0x0x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim32_causal0x0x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim32_causal0x1x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim32_causal0x1x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim32_causal1x0x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim32_causal1x0x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim32_causal1x1x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim32_causal1x1x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim64(is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim64_causal0x0x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim64_causal0x0x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim64_causal0x1x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim64_causal0x1x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim64_causal1x0x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim64_causal1x0x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim64_causal1x1x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim64_causal1x1x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim128(is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim128_causal0x0x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim128_causal0x0x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim128_causal0x1x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim128_causal0x1x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim128_causal1x0x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim128_causal1x0x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim128_causal1x1x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim128_causal1x1x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim256(is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim256_causal0x0x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim256_causal0x0x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim256_causal0x1x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim256_causal0x1x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim256_causal1x0x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim256_causal1x0x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim256_causal1x1x0(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim256_causal1x1x1(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 256) got invalid causal mask!"); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fna_hopper +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_dtype.h b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_dtype.h new file mode 100644 index 00000000..385551f6 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_dtype.h @@ -0,0 +1,69 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_hopper { +#define DISPATCH_HOPPER_FNA_FORWARD_1D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (dtype == torch::kFloat16) { \ + DISPATCH_HOPPER_FNA_FORWARD_1D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (dtype == torch::kBFloat16) { \ + DISPATCH_HOPPER_FNA_FORWARD_1D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-1D forward kernel dispatch failed! It does not support dtype " + std::string(c10::toString(dtype)) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (dtype == torch::kFloat16) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (dtype == torch::kBFloat16) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It does not support dtype " + std::string(c10::toString(dtype)) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (dtype == torch::kFloat16) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (dtype == torch::kBFloat16) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It does not support dtype " + std::string(c10::toString(dtype)) + "."); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fna_hopper +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_head_dim.h b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_head_dim.h new file mode 100644 index 00000000..8569c6c3 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_head_dim.h @@ -0,0 +1,144 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_hopper { +#define DISPATCH_HOPPER_FNA_FORWARD_1D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (dim == 32) { \ + DISPATCH_HOPPER_FNA_FORWARD_1D_float16_headdim32(is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (dim == 64) { \ + DISPATCH_HOPPER_FNA_FORWARD_1D_float16_headdim64(is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (dim == 128) { \ + DISPATCH_HOPPER_FNA_FORWARD_1D_float16_headdim128(is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (dim == 256) { \ + DISPATCH_HOPPER_FNA_FORWARD_1D_float16_headdim256(is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA forward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for float16."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_1D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (dim == 32) { \ + DISPATCH_HOPPER_FNA_FORWARD_1D_bfloat16_headdim32(is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (dim == 64) { \ + DISPATCH_HOPPER_FNA_FORWARD_1D_bfloat16_headdim64(is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (dim == 128) { \ + DISPATCH_HOPPER_FNA_FORWARD_1D_bfloat16_headdim128(is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (dim == 256) { \ + DISPATCH_HOPPER_FNA_FORWARD_1D_bfloat16_headdim256(is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA forward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for bfloat16."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (dim == 32) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim32(is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (dim == 64) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim64(is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (dim == 128) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim128(is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (dim == 256) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim256(is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA forward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for float16."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (dim == 32) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim32(is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (dim == 64) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim64(is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (dim == 128) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim128(is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (dim == 256) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim256(is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA forward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for bfloat16."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (dim == 32) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim32(is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (dim == 64) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim64(is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (dim == 128) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim128(is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (dim == 256) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim256(is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA forward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for float16."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (dim == 32) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim32(is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (dim == 64) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim64(is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (dim == 128) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim128(is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if (dim == 256) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim256(is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA forward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for bfloat16."); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fna_hopper +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_tile_shape.h b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_tile_shape.h new file mode 100644 index 00000000..c28c1ebd --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_tile_shape.h @@ -0,0 +1,2034 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_hopper { +#define DISPATCH_HOPPER_FNA_FORWARD_1D_float16_headdim32_causal0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 64 && \ +cute::get<0>(kv_tile_shape) == 128 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna1d_float16_64x128x32_Q64_KV128_causal0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-1D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_1D_float16_headdim32_causal1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 64 && \ +cute::get<0>(kv_tile_shape) == 128 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna1d_float16_64x128x32_Q64_KV128_causal1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-1D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_1D_float16_headdim64_causal0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 64 && \ +cute::get<0>(kv_tile_shape) == 128 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna1d_float16_64x128x64_Q64_KV128_causal0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-1D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_1D_float16_headdim64_causal1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 64 && \ +cute::get<0>(kv_tile_shape) == 128 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna1d_float16_64x128x64_Q64_KV128_causal1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-1D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_1D_float16_headdim128_causal0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 128 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna1d_float16_128x128x128_coop_Q128_KV128_causal0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 128 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna1d_float16_128x128x128_pp_Q128_KV128_causal0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-1D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_1D_float16_headdim128_causal1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 128 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna1d_float16_128x128x128_coop_Q128_KV128_causal1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 128 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna1d_float16_128x128x128_pp_Q128_KV128_causal1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-1D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_1D_float16_headdim256_causal0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 64 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna1d_float16_128x64x256_coop_Q128_KV64_causal0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-1D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 256): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_1D_float16_headdim256_causal1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 64 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna1d_float16_128x64x256_coop_Q128_KV64_causal1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-1D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 256): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_1D_bfloat16_headdim32_causal0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 64 && \ +cute::get<0>(kv_tile_shape) == 128 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna1d_bfloat16_64x128x32_Q64_KV128_causal0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-1D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_1D_bfloat16_headdim32_causal1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 64 && \ +cute::get<0>(kv_tile_shape) == 128 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna1d_bfloat16_64x128x32_Q64_KV128_causal1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-1D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_1D_bfloat16_headdim64_causal0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 64 && \ +cute::get<0>(kv_tile_shape) == 128 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna1d_bfloat16_64x128x64_Q64_KV128_causal0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-1D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_1D_bfloat16_headdim64_causal1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 64 && \ +cute::get<0>(kv_tile_shape) == 128 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna1d_bfloat16_64x128x64_Q64_KV128_causal1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-1D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_1D_bfloat16_headdim128_causal0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 128 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna1d_bfloat16_128x128x128_coop_Q128_KV128_causal0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 128 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna1d_bfloat16_128x128x128_pp_Q128_KV128_causal0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-1D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_1D_bfloat16_headdim128_causal1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 128 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna1d_bfloat16_128x128x128_coop_Q128_KV128_causal1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 128 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna1d_bfloat16_128x128x128_pp_Q128_KV128_causal1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-1D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_1D_bfloat16_headdim256_causal0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 64 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna1d_bfloat16_128x64x256_coop_Q128_KV64_causal0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-1D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 256): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_1D_bfloat16_headdim256_causal1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 64 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna1d_bfloat16_128x64x256_coop_Q128_KV64_causal1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-1D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 256): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim32_causal0x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_64x128x32_Q8x8_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_64x128x32_Q8x8_KV8x16_causal0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim32_causal0x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_64x128x32_Q8x8_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_64x128x32_Q8x8_KV8x16_causal0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim32_causal1x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_64x128x32_Q8x8_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_64x128x32_Q8x8_KV8x16_causal1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim32_causal1x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_64x128x32_Q8x8_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_64x128x32_Q8x8_KV8x16_causal1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim64_causal0x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_64x128x64_Q8x8_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_64x128x64_Q8x8_KV8x16_causal0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim64_causal0x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_64x128x64_Q8x8_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_64x128x64_Q8x8_KV8x16_causal0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim64_causal1x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_64x128x64_Q8x8_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_64x128x64_Q8x8_KV8x16_causal1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim64_causal1x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_64x128x64_Q8x8_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_64x128x64_Q8x8_KV8x16_causal1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim128_causal0x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_128x128x128_coop_Q16x8_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_128x128x128_pp_Q16x8_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim128_causal0x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_128x128x128_coop_Q16x8_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_128x128x128_pp_Q16x8_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim128_causal1x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_128x128x128_coop_Q16x8_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_128x128x128_pp_Q16x8_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim128_causal1x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_128x128x128_coop_Q16x8_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_128x128x128_pp_Q16x8_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim256_causal0x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_128x64x256_coop_Q16x8_KV8x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_128x64x256_coop_Q8x16_KV8x8_causal0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 256): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim256_causal0x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_128x64x256_coop_Q16x8_KV8x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_128x64x256_coop_Q8x16_KV8x8_causal0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 256): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim256_causal1x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_128x64x256_coop_Q16x8_KV8x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_128x64x256_coop_Q8x16_KV8x8_causal1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 256): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_float16_headdim256_causal1x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_128x64x256_coop_Q16x8_KV8x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna2d_float16_128x64x256_coop_Q8x16_KV8x8_causal1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 256): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim32_causal0x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_64x128x32_Q8x8_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_64x128x32_Q8x8_KV8x16_causal0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim32_causal0x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_64x128x32_Q8x8_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_64x128x32_Q8x8_KV8x16_causal0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim32_causal1x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_64x128x32_Q8x8_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_64x128x32_Q8x8_KV8x16_causal1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim32_causal1x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_64x128x32_Q8x8_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_64x128x32_Q8x8_KV8x16_causal1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim64_causal0x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_64x128x64_Q8x8_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_64x128x64_Q8x8_KV8x16_causal0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim64_causal0x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_64x128x64_Q8x8_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_64x128x64_Q8x8_KV8x16_causal0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim64_causal1x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_64x128x64_Q8x8_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_64x128x64_Q8x8_KV8x16_causal1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim64_causal1x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_64x128x64_Q8x8_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_64x128x64_Q8x8_KV8x16_causal1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim128_causal0x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_128x128x128_coop_Q16x8_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_128x128x128_pp_Q16x8_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim128_causal0x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_128x128x128_coop_Q16x8_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_128x128x128_pp_Q16x8_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim128_causal1x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_128x128x128_coop_Q16x8_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_128x128x128_pp_Q16x8_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim128_causal1x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_128x128x128_coop_Q16x8_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_128x128x128_pp_Q16x8_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim256_causal0x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_128x64x256_coop_Q16x8_KV8x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_128x64x256_coop_Q8x16_KV8x8_causal0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 256): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim256_causal0x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_128x64x256_coop_Q16x8_KV8x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_128x64x256_coop_Q8x16_KV8x8_causal0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 256): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim256_causal1x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_128x64x256_coop_Q16x8_KV8x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_128x64x256_coop_Q8x16_KV8x8_causal1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 256): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16_headdim256_causal1x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_128x64x256_coop_Q16x8_KV8x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 16 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna2d_bfloat16_128x64x256_coop_Q8x16_KV8x8_causal1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 256): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim32_causal0x0x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x32_Q4x4x4_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x32_Q4x4x4_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim32_causal0x0x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x32_Q4x4x4_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x32_Q4x4x4_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim32_causal0x1x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x32_Q4x4x4_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x32_Q4x4x4_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim32_causal0x1x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x32_Q4x4x4_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x32_Q4x4x4_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim32_causal1x0x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x32_Q4x4x4_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x32_Q4x4x4_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim32_causal1x0x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x32_Q4x4x4_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x32_Q4x4x4_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim32_causal1x1x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x32_Q4x4x4_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x32_Q4x4x4_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim32_causal1x1x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x32_Q4x4x4_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x32_Q4x4x4_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim64_causal0x0x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x64_Q4x4x4_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x64_Q4x4x4_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim64_causal0x0x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x64_Q4x4x4_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x64_Q4x4x4_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim64_causal0x1x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x64_Q4x4x4_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x64_Q4x4x4_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim64_causal0x1x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x64_Q4x4x4_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x64_Q4x4x4_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim64_causal1x0x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x64_Q4x4x4_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x64_Q4x4x4_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim64_causal1x0x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x64_Q4x4x4_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x64_Q4x4x4_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim64_causal1x1x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x64_Q4x4x4_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x64_Q4x4x4_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim64_causal1x1x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x64_Q4x4x4_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_64x128x64_Q4x4x4_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim128_causal0x0x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_coop_Q4x4x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_pp_Q4x4x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_coop_Q2x8x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_pp_Q2x8x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim128_causal0x0x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_coop_Q4x4x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_pp_Q4x4x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_coop_Q2x8x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_pp_Q2x8x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim128_causal0x1x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_coop_Q4x4x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_pp_Q4x4x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_coop_Q2x8x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_pp_Q2x8x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim128_causal0x1x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_coop_Q4x4x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_pp_Q4x4x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_coop_Q2x8x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_pp_Q2x8x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim128_causal1x0x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_coop_Q4x4x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_pp_Q4x4x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_coop_Q2x8x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_pp_Q2x8x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim128_causal1x0x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_coop_Q4x4x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_pp_Q4x4x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_coop_Q2x8x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_pp_Q2x8x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim128_causal1x1x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_coop_Q4x4x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_pp_Q4x4x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_coop_Q2x8x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_pp_Q2x8x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim128_causal1x1x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_coop_Q4x4x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_pp_Q4x4x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_coop_Q2x8x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x128x128_pp_Q2x8x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim256_causal0x0x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x64x256_coop_Q4x4x8_KV4x4x4_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x64x256_coop_Q2x8x8_KV4x4x4_causal0x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 256): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim256_causal0x0x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x64x256_coop_Q4x4x8_KV4x4x4_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x64x256_coop_Q2x8x8_KV4x4x4_causal0x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 256): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim256_causal0x1x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x64x256_coop_Q4x4x8_KV4x4x4_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x64x256_coop_Q2x8x8_KV4x4x4_causal0x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 256): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim256_causal0x1x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x64x256_coop_Q4x4x8_KV4x4x4_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x64x256_coop_Q2x8x8_KV4x4x4_causal0x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 256): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim256_causal1x0x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x64x256_coop_Q4x4x8_KV4x4x4_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x64x256_coop_Q2x8x8_KV4x4x4_causal1x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 256): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim256_causal1x0x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x64x256_coop_Q4x4x8_KV4x4x4_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x64x256_coop_Q2x8x8_KV4x4x4_causal1x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 256): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim256_causal1x1x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x64x256_coop_Q4x4x8_KV4x4x4_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x64x256_coop_Q2x8x8_KV4x4x4_causal1x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 256): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_float16_headdim256_causal1x1x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x64x256_coop_Q4x4x8_KV4x4x4_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_float16_128x64x256_coop_Q2x8x8_KV4x4x4_causal1x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (float16, head_dim 256): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim32_causal0x0x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim32_causal0x0x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim32_causal0x1x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim32_causal0x1x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim32_causal1x0x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim32_causal1x0x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim32_causal1x1x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim32_causal1x1x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim64_causal0x0x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim64_causal0x0x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim64_causal0x1x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim64_causal0x1x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim64_causal1x0x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim64_causal1x0x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim64_causal1x1x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim64_causal1x1x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::NonPersistent) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim128_causal0x0x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_coop_Q4x4x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_pp_Q4x4x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_coop_Q2x8x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_pp_Q2x8x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim128_causal0x0x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_coop_Q4x4x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_pp_Q4x4x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_coop_Q2x8x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_pp_Q2x8x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim128_causal0x1x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_coop_Q4x4x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_pp_Q4x4x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_coop_Q2x8x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_pp_Q2x8x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim128_causal0x1x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_coop_Q4x4x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_pp_Q4x4x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_coop_Q2x8x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_pp_Q2x8x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim128_causal1x0x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_coop_Q4x4x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_pp_Q4x4x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_coop_Q2x8x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_pp_Q2x8x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim128_causal1x0x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_coop_Q4x4x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_pp_Q4x4x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_coop_Q2x8x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_pp_Q2x8x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim128_causal1x1x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_coop_Q4x4x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_pp_Q4x4x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_coop_Q2x8x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_pp_Q2x8x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim128_causal1x1x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_coop_Q4x4x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_pp_Q4x4x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_coop_Q2x8x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSPingpong) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x128x128_pp_Q2x8x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim256_causal0x0x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x64x256_coop_Q4x4x8_KV4x4x4_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x64x256_coop_Q2x8x8_KV4x4x4_causal0x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 256): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim256_causal0x0x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x64x256_coop_Q4x4x8_KV4x4x4_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x64x256_coop_Q2x8x8_KV4x4x4_causal0x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 256): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim256_causal0x1x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x64x256_coop_Q4x4x8_KV4x4x4_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x64x256_coop_Q2x8x8_KV4x4x4_causal0x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 256): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim256_causal0x1x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x64x256_coop_Q4x4x8_KV4x4x4_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x64x256_coop_Q2x8x8_KV4x4x4_causal0x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 256): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim256_causal1x0x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x64x256_coop_Q4x4x8_KV4x4x4_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x64x256_coop_Q2x8x8_KV4x4x4_causal1x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 256): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim256_causal1x0x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x64x256_coop_Q4x4x8_KV4x4x4_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x64x256_coop_Q2x8x8_KV4x4x4_causal1x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 256): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim256_causal1x1x0(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x64x256_coop_Q4x4x8_KV4x4x4_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x64x256_coop_Q2x8x8_KV4x4x4_causal1x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 256): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16_headdim256_causal1x1x1(q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x64x256_coop_Q4x4x8_KV4x4x4_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 4 && \ +kernel_type == natten::cuda::hopper::HopperKernelSchedule::WSCooperative) { \ + natten::cuda::fna_hopper::hopper_fna3d_bfloat16_128x64x256_coop_Q2x8x8_KV4x4x4_causal1x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D forward kernel dispatch failed! It got invalid Q tile, KV tile, and schedule combination (bfloat16, head_dim 256): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + "), schedule=" + natten::cuda::hopper::to_string(kernel_type) + "."); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fna_hopper +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna/interface.h b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna/interface.h new file mode 100644 index 00000000..2d0d3145 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna/interface.h @@ -0,0 +1,46 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_hopper { +#define DISPATCH_HOPPER_FNA_FORWARD(rank, dtype, dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ + [&] { \ + if constexpr (rank == 1) { \ + DISPATCH_HOPPER_FNA_FORWARD_1D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if constexpr (rank == 2) { \ + DISPATCH_HOPPER_FNA_FORWARD_2D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else if constexpr (rank == 3) { \ + DISPATCH_HOPPER_FNA_FORWARD_3D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA forward kernel dispatch failed! It only supports NA1D, 2D, and 3D!"); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fna_hopper +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna/kernels.h b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna/kernels.h new file mode 100644 index 00000000..047b3eca --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna/kernels.h @@ -0,0 +1,5397 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_hopper { + +void hopper_fna1d_float16_64x128x32_Q64_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_float16_64x128x32_Q64_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_float16_64x128x64_Q64_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_float16_64x128x64_Q64_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_float16_128x128x128_coop_Q128_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_float16_128x128x128_pp_Q128_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_float16_128x128x128_coop_Q128_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_float16_128x128x128_pp_Q128_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_float16_128x64x256_coop_Q128_KV64_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_float16_128x64x256_coop_Q128_KV64_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_bfloat16_64x128x32_Q64_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_bfloat16_64x128x32_Q64_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_bfloat16_64x128x64_Q64_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_bfloat16_64x128x64_Q64_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_bfloat16_128x128x128_coop_Q128_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_bfloat16_128x128x128_pp_Q128_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_bfloat16_128x128x128_coop_Q128_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_bfloat16_128x128x128_pp_Q128_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_bfloat16_128x64x256_coop_Q128_KV64_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_bfloat16_128x64x256_coop_Q128_KV64_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_64x128x32_Q8x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_64x128x32_Q8x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_64x128x32_Q8x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_64x128x32_Q8x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_64x128x32_Q8x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_64x128x32_Q8x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_64x128x32_Q8x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_64x128x32_Q8x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_64x128x64_Q8x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_64x128x64_Q8x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_64x128x64_Q8x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_64x128x64_Q8x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_64x128x64_Q8x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_64x128x64_Q8x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_64x128x64_Q8x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_64x128x64_Q8x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_128x128x128_coop_Q16x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_128x128x128_pp_Q16x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_128x128x128_coop_Q16x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_128x128x128_pp_Q16x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_128x128x128_coop_Q16x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_128x128x128_pp_Q16x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_128x128x128_coop_Q16x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_128x128x128_pp_Q16x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_128x64x256_coop_Q16x8_KV8x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_128x64x256_coop_Q8x16_KV8x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_128x64x256_coop_Q16x8_KV8x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_128x64x256_coop_Q8x16_KV8x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_128x64x256_coop_Q16x8_KV8x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_128x64x256_coop_Q8x16_KV8x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_128x64x256_coop_Q16x8_KV8x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_float16_128x64x256_coop_Q8x16_KV8x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_64x128x32_Q8x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_64x128x32_Q8x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_64x128x32_Q8x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_64x128x32_Q8x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_64x128x32_Q8x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_64x128x32_Q8x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_64x128x32_Q8x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_64x128x32_Q8x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_64x128x64_Q8x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_64x128x64_Q8x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_64x128x64_Q8x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_64x128x64_Q8x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_64x128x64_Q8x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_64x128x64_Q8x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_64x128x64_Q8x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_64x128x64_Q8x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_128x128x128_coop_Q16x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_128x128x128_pp_Q16x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_128x128x128_coop_Q16x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_128x128x128_pp_Q16x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_128x128x128_coop_Q16x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_128x128x128_pp_Q16x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_128x128x128_coop_Q16x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_128x128x128_pp_Q16x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_128x64x256_coop_Q16x8_KV8x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_128x64x256_coop_Q8x16_KV8x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_128x64x256_coop_Q16x8_KV8x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_128x64x256_coop_Q8x16_KV8x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_128x64x256_coop_Q16x8_KV8x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_128x64x256_coop_Q8x16_KV8x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_128x64x256_coop_Q16x8_KV8x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_bfloat16_128x64x256_coop_Q8x16_KV8x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_coop_Q4x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_pp_Q4x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_coop_Q2x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_pp_Q2x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_coop_Q4x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_pp_Q4x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_coop_Q2x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_pp_Q2x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_coop_Q4x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_pp_Q4x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_coop_Q2x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_pp_Q2x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_coop_Q4x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_pp_Q4x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_coop_Q2x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_pp_Q2x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_coop_Q4x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_pp_Q4x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_coop_Q2x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_pp_Q2x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_coop_Q4x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_pp_Q4x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_coop_Q2x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_pp_Q2x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_coop_Q4x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_pp_Q4x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_coop_Q2x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_pp_Q2x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_coop_Q4x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_pp_Q4x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_coop_Q2x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x128x128_pp_Q2x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x64x256_coop_Q4x4x8_KV4x4x4_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x64x256_coop_Q2x8x8_KV4x4x4_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x64x256_coop_Q4x4x8_KV4x4x4_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x64x256_coop_Q2x8x8_KV4x4x4_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x64x256_coop_Q4x4x8_KV4x4x4_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x64x256_coop_Q2x8x8_KV4x4x4_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x64x256_coop_Q4x4x8_KV4x4x4_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x64x256_coop_Q2x8x8_KV4x4x4_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x64x256_coop_Q4x4x8_KV4x4x4_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x64x256_coop_Q2x8x8_KV4x4x4_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x64x256_coop_Q4x4x8_KV4x4x4_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x64x256_coop_Q2x8x8_KV4x4x4_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x64x256_coop_Q4x4x8_KV4x4x4_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x64x256_coop_Q2x8x8_KV4x4x4_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x64x256_coop_Q4x4x8_KV4x4x4_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_float16_128x64x256_coop_Q2x8x8_KV4x4x4_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_coop_Q4x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_pp_Q4x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_coop_Q2x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_pp_Q2x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_coop_Q4x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_pp_Q4x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_coop_Q2x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_pp_Q2x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_coop_Q4x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_pp_Q4x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_coop_Q2x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_pp_Q2x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_coop_Q4x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_pp_Q4x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_coop_Q2x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_pp_Q2x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_coop_Q4x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_pp_Q4x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_coop_Q2x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_pp_Q2x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_coop_Q4x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_pp_Q4x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_coop_Q2x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_pp_Q2x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_coop_Q4x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_pp_Q4x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_coop_Q2x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_pp_Q2x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_coop_Q4x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_pp_Q4x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_coop_Q2x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x128x128_pp_Q2x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x64x256_coop_Q4x4x8_KV4x4x4_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x64x256_coop_Q2x8x8_KV4x4x4_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x64x256_coop_Q4x4x8_KV4x4x4_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x64x256_coop_Q2x8x8_KV4x4x4_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x64x256_coop_Q4x4x8_KV4x4x4_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x64x256_coop_Q2x8x8_KV4x4x4_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x64x256_coop_Q4x4x8_KV4x4x4_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x64x256_coop_Q2x8x8_KV4x4x4_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x64x256_coop_Q4x4x8_KV4x4x4_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x64x256_coop_Q2x8x8_KV4x4x4_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x64x256_coop_Q4x4x8_KV4x4x4_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x64x256_coop_Q2x8x8_KV4x4x4_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x64x256_coop_Q4x4x8_KV4x4x4_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x64x256_coop_Q2x8x8_KV4x4x4_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x64x256_coop_Q4x4x8_KV4x4x4_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_bfloat16_128x64x256_coop_Q2x8x8_KV4x4x4_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + + +} // namespace natten +} // namespace cuda +} // namespace fna_hopper +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_cm.h b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_cm.h new file mode 100644 index 00000000..327561ae --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_cm.h @@ -0,0 +1,407 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_hopper { +#define DISPATCH_HOPPER_FNA_BACKWARD_1D_float16_headdim32(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_1D_float16_headdim32_causal0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_1D_float16_headdim32_causal1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-1D backward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_1D_float16_headdim64(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_1D_float16_headdim64_causal0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_1D_float16_headdim64_causal1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-1D backward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_1D_float16_headdim128(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_1D_float16_headdim128_causal0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_1D_float16_headdim128_causal1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-1D backward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_1D_bfloat16_headdim32(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_1D_bfloat16_headdim32_causal0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_1D_bfloat16_headdim32_causal1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-1D backward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_1D_bfloat16_headdim64(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_1D_bfloat16_headdim64_causal0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_1D_bfloat16_headdim64_causal1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-1D backward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_1D_bfloat16_headdim128(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_1D_bfloat16_headdim128_causal0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_1D_bfloat16_headdim128_causal1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-1D backward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim32(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim32_causal0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim32_causal0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim32_causal1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim32_causal1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim64(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim64_causal0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim64_causal0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim64_causal1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim64_causal1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim128(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim128_causal0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim128_causal0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim128_causal1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim128_causal1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim32(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim32_causal0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim32_causal0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim32_causal1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim32_causal1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim64(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim64_causal0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim64_causal0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim64_causal1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim64_causal1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim128(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim128_causal0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim128_causal0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim128_causal1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim128_causal1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim32(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim32_causal0x0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim32_causal0x0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim32_causal0x1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim32_causal0x1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim32_causal1x0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim32_causal1x0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim32_causal1x1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim32_causal1x1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim64(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim64_causal0x0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim64_causal0x0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim64_causal0x1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim64_causal0x1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim64_causal1x0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim64_causal1x0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim64_causal1x1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim64_causal1x1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim128(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim128_causal0x0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim128_causal0x0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim128_causal0x1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim128_causal0x1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim128_causal1x0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim128_causal1x0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim128_causal1x1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim128_causal1x1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! Causal mask dispatcher (float16, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim32(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim32_causal0x0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim32_causal0x0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim32_causal0x1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim32_causal0x1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim32_causal1x0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim32_causal1x0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim32_causal1x1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim32_causal1x1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 32) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim64(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim64_causal0x0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim64_causal0x0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim64_causal0x1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim64_causal0x1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim64_causal1x0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim64_causal1x0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim64_causal1x1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim64_causal1x1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 64) got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim128(is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim128_causal0x0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim128_causal0x0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim128_causal0x1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim128_causal0x1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim128_causal1x0x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim128_causal1x0x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim128_causal1x1x0(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim128_causal1x1x1(q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! Causal mask dispatcher (bfloat16, head_dim 128) got invalid causal mask!"); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fna_hopper +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_dtype.h b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_dtype.h new file mode 100644 index 00000000..04ff1e89 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_dtype.h @@ -0,0 +1,68 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_hopper { +#define DISPATCH_HOPPER_FNA_BACKWARD_1D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (dtype == torch::kFloat16) { \ + DISPATCH_HOPPER_FNA_BACKWARD_1D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dtype == torch::kBFloat16) { \ + DISPATCH_HOPPER_FNA_BACKWARD_1D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-1D backward kernel dispatch failed! It does not support dtype " + std::string(c10::toString(dtype)) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (dtype == torch::kFloat16) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dtype == torch::kBFloat16) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! It does not support dtype " + std::string(c10::toString(dtype)) + "."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (dtype == torch::kFloat16) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dtype == torch::kBFloat16) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It does not support dtype " + std::string(c10::toString(dtype)) + "."); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fna_hopper +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_head_dim.h b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_head_dim.h new file mode 100644 index 00000000..fd5e8c19 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_head_dim.h @@ -0,0 +1,125 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_hopper { +#define DISPATCH_HOPPER_FNA_BACKWARD_1D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (dim == 32) { \ + DISPATCH_HOPPER_FNA_BACKWARD_1D_float16_headdim32(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dim == 64) { \ + DISPATCH_HOPPER_FNA_BACKWARD_1D_float16_headdim64(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dim == 128) { \ + DISPATCH_HOPPER_FNA_BACKWARD_1D_float16_headdim128(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA backward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for float16."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_1D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (dim == 32) { \ + DISPATCH_HOPPER_FNA_BACKWARD_1D_bfloat16_headdim32(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dim == 64) { \ + DISPATCH_HOPPER_FNA_BACKWARD_1D_bfloat16_headdim64(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dim == 128) { \ + DISPATCH_HOPPER_FNA_BACKWARD_1D_bfloat16_headdim128(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA backward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for bfloat16."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (dim == 32) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim32(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dim == 64) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim64(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dim == 128) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim128(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA backward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for float16."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (dim == 32) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim32(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dim == 64) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim64(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dim == 128) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim128(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA backward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for bfloat16."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (dim == 32) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim32(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dim == 64) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim64(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dim == 128) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim128(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA backward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for float16."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (dim == 32) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim32(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dim == 64) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim64(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if (dim == 128) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim128(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA backward kernel dispatch failed! It does not support head dim " + std::to_string(dim) + " for bfloat16."); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fna_hopper +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_tile_shape.h b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_tile_shape.h new file mode 100644 index 00000000..e1588500 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_tile_shape.h @@ -0,0 +1,1785 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_hopper { +#define DISPATCH_HOPPER_FNA_BACKWARD_1D_float16_headdim32_causal0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 64 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_hopper::hopper_fna1d_backward_float16_64x128x32_Q64_KV128_causal0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_hopper::hopper_fna1d_backward_float16_128x128x32_Q128_KV128_causal0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-1D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_1D_float16_headdim32_causal1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 64 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_hopper::hopper_fna1d_backward_float16_64x128x32_Q64_KV128_causal1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_hopper::hopper_fna1d_backward_float16_128x128x32_Q128_KV128_causal1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-1D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_1D_float16_headdim64_causal0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 64 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_hopper::hopper_fna1d_backward_float16_64x128x64_Q64_KV128_causal0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_hopper::hopper_fna1d_backward_float16_128x128x64_Q128_KV128_causal0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-1D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_1D_float16_headdim64_causal1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 64 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_hopper::hopper_fna1d_backward_float16_64x128x64_Q64_KV128_causal1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_hopper::hopper_fna1d_backward_float16_128x128x64_Q128_KV128_causal1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-1D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_1D_float16_headdim128_causal0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 64 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_hopper::hopper_fna1d_backward_float16_64x128x128_Q64_KV128_causal0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-1D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_1D_float16_headdim128_causal1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 64 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_hopper::hopper_fna1d_backward_float16_64x128x128_Q64_KV128_causal1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-1D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_1D_bfloat16_headdim32_causal0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 64 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_hopper::hopper_fna1d_backward_bfloat16_64x128x32_Q64_KV128_causal0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_hopper::hopper_fna1d_backward_bfloat16_128x128x32_Q128_KV128_causal0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-1D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_1D_bfloat16_headdim32_causal1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 64 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_hopper::hopper_fna1d_backward_bfloat16_64x128x32_Q64_KV128_causal1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_hopper::hopper_fna1d_backward_bfloat16_128x128x32_Q128_KV128_causal1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-1D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_1D_bfloat16_headdim64_causal0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 64 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_hopper::hopper_fna1d_backward_bfloat16_64x128x64_Q64_KV128_causal0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_hopper::hopper_fna1d_backward_bfloat16_128x128x64_Q128_KV128_causal0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-1D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_1D_bfloat16_headdim64_causal1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 64 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_hopper::hopper_fna1d_backward_bfloat16_64x128x64_Q64_KV128_causal1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 128 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_hopper::hopper_fna1d_backward_bfloat16_128x128x64_Q128_KV128_causal1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-1D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_1D_bfloat16_headdim128_causal0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 64 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_hopper::hopper_fna1d_backward_bfloat16_64x128x128_Q64_KV128_causal0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-1D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_1D_bfloat16_headdim128_causal1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 64 && \ +cute::get<0>(kv_tile_shape) == 128) { \ + natten::cuda::fna_hopper::hopper_fna1d_backward_bfloat16_64x128x128_Q64_KV128_causal1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-1D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim32_causal0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_64x128x32_Q8x8_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_64x128x32_Q8x8_KV8x16_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_128x128x32_Q16x8_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_128x128x32_Q16x8_KV8x16_causal0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim32_causal0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_64x128x32_Q8x8_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_64x128x32_Q8x8_KV8x16_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_128x128x32_Q16x8_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_128x128x32_Q16x8_KV8x16_causal0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim32_causal1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_64x128x32_Q8x8_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_64x128x32_Q8x8_KV8x16_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_128x128x32_Q16x8_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_128x128x32_Q16x8_KV8x16_causal1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim32_causal1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_64x128x32_Q8x8_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_64x128x32_Q8x8_KV8x16_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_128x128x32_Q16x8_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_128x128x32_Q16x8_KV8x16_causal1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim64_causal0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_64x128x64_Q8x8_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_64x128x64_Q8x8_KV8x16_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_128x128x64_Q16x8_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_128x128x64_Q16x8_KV8x16_causal0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim64_causal0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_64x128x64_Q8x8_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_64x128x64_Q8x8_KV8x16_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_128x128x64_Q16x8_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_128x128x64_Q16x8_KV8x16_causal0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim64_causal1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_64x128x64_Q8x8_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_64x128x64_Q8x8_KV8x16_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_128x128x64_Q16x8_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_128x128x64_Q16x8_KV8x16_causal1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim64_causal1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_64x128x64_Q8x8_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_64x128x64_Q8x8_KV8x16_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_128x128x64_Q16x8_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_128x128x64_Q16x8_KV8x16_causal1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim128_causal0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_64x128x128_Q8x8_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_64x128x128_Q8x8_KV8x16_causal0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim128_causal0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_64x128x128_Q8x8_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_64x128x128_Q8x8_KV8x16_causal0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim128_causal1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_64x128x128_Q8x8_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_64x128x128_Q8x8_KV8x16_causal1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_float16_headdim128_causal1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_64x128x128_Q8x8_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_float16_64x128x128_Q8x8_KV8x16_causal1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim32_causal0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_64x128x32_Q8x8_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_64x128x32_Q8x8_KV8x16_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_128x128x32_Q16x8_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_128x128x32_Q16x8_KV8x16_causal0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim32_causal0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_64x128x32_Q8x8_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_64x128x32_Q8x8_KV8x16_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_128x128x32_Q16x8_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_128x128x32_Q16x8_KV8x16_causal0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim32_causal1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_64x128x32_Q8x8_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_64x128x32_Q8x8_KV8x16_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_128x128x32_Q16x8_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_128x128x32_Q16x8_KV8x16_causal1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim32_causal1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_64x128x32_Q8x8_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_64x128x32_Q8x8_KV8x16_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_128x128x32_Q16x8_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_128x128x32_Q16x8_KV8x16_causal1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim64_causal0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_64x128x64_Q8x8_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_64x128x64_Q8x8_KV8x16_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_128x128x64_Q16x8_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_128x128x64_Q16x8_KV8x16_causal0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim64_causal0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_64x128x64_Q8x8_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_64x128x64_Q8x8_KV8x16_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_128x128x64_Q16x8_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_128x128x64_Q16x8_KV8x16_causal0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim64_causal1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_64x128x64_Q8x8_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_64x128x64_Q8x8_KV8x16_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_128x128x64_Q16x8_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_128x128x64_Q16x8_KV8x16_causal1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim64_causal1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_64x128x64_Q8x8_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_64x128x64_Q8x8_KV8x16_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_128x128x64_Q16x8_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 16 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_128x128x64_Q16x8_KV8x16_causal1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim128_causal0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_64x128x128_Q8x8_KV16x8_causal0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_64x128x128_Q8x8_KV8x16_causal0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim128_causal0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_64x128x128_Q8x8_KV16x8_causal0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_64x128x128_Q8x8_KV8x16_causal0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim128_causal1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_64x128x128_Q8x8_KV16x8_causal1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_64x128x128_Q8x8_KV8x16_causal1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16_headdim128_causal1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 16 && cute::get<1>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_64x128x128_Q8x8_KV16x8_causal1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 8 && cute::get<1>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 8 && cute::get<1>(kv_tile_shape) == 16) { \ + natten::cuda::fna_hopper::hopper_fna2d_backward_bfloat16_64x128x128_Q8x8_KV8x16_causal1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-2D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim32_causal0x0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim32_causal0x0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim32_causal0x1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim32_causal0x1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim32_causal1x0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim32_causal1x0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim32_causal1x1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim32_causal1x1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim64_causal0x0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim64_causal0x0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim64_causal0x1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim64_causal0x1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim64_causal1x0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim64_causal1x0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim64_causal1x1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim64_causal1x1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim128_causal0x0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q2x4x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q1x8x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim128_causal0x0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q2x4x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q1x8x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim128_causal0x1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q2x4x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q1x8x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim128_causal0x1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q2x4x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q1x8x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim128_causal1x0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q2x4x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q1x8x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim128_causal1x0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q2x4x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q1x8x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim128_causal1x1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q2x4x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q1x8x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_float16_headdim128_causal1x1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q2x4x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_float16_64x128x128_Q1x8x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (float16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim32_causal0x0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim32_causal0x0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim32_causal0x1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim32_causal0x1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim32_causal1x0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim32_causal1x0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim32_causal1x1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim32_causal1x1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 32): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim64_causal0x0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim64_causal0x0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim64_causal0x1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim64_causal0x1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim64_causal1x0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim64_causal1x0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim64_causal1x1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim64_causal1x1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 64): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim128_causal0x0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV4x4x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q2x4x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q1x8x8_KV2x8x8_causal0x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim128_causal0x0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV4x4x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q2x4x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q1x8x8_KV2x8x8_causal0x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim128_causal0x1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV4x4x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q2x4x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q1x8x8_KV2x8x8_causal0x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim128_causal0x1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV4x4x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q2x4x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q1x8x8_KV2x8x8_causal0x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim128_causal1x0x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV4x4x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q2x4x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q1x8x8_KV2x8x8_causal1x0x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim128_causal1x0x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV4x4x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q2x4x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q1x8x8_KV2x8x8_causal1x0x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim128_causal1x1x0(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV4x4x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q2x4x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q1x8x8_KV2x8x8_causal1x1x0(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + +#define DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16_headdim128_causal1x1x1(q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 4 && cute::get<1>(kv_tile_shape) == 4 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV4x4x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 4 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 4 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 2 && cute::get<1>(q_tile_shape) == 4 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q2x4x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else if (cute::get<0>(q_tile_shape) == 1 && cute::get<1>(q_tile_shape) == 8 && cute::get<2>(q_tile_shape) == 8 && \ +cute::get<0>(kv_tile_shape) == 2 && cute::get<1>(kv_tile_shape) == 8 && cute::get<2>(kv_tile_shape) == 8) { \ + natten::cuda::fna_hopper::hopper_fna3d_backward_bfloat16_64x128x128_Q1x8x8_KV2x8x8_causal1x1x1(__VA_ARGS__); \ +} \ + else { \ + throw std::runtime_error("Hopper FNA-3D backward kernel dispatch failed! It got invalid Q tile and KV tile combination (bfloat16, head_dim 128): q_tile=(" + std::to_string(cute::get<0>(q_tile_shape)) + "," + std::to_string(cute::get<1>(q_tile_shape)) + "," + std::to_string(cute::get<2>(q_tile_shape)) + "), kv_tile=(" + std::to_string(cute::get<0>(kv_tile_shape)) + "," + std::to_string(cute::get<1>(kv_tile_shape)) + "," + std::to_string(cute::get<2>(kv_tile_shape)) + ")."); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fna_hopper +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/interface.h b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/interface.h new file mode 100644 index 00000000..45720df4 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/interface.h @@ -0,0 +1,45 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_hopper { +#define DISPATCH_HOPPER_FNA_BACKWARD(rank, dtype, dim, is_causal, q_tile_shape, kv_tile_shape, ...) \ + [&] { \ + if constexpr (rank == 1) { \ + DISPATCH_HOPPER_FNA_BACKWARD_1D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if constexpr (rank == 2) { \ + DISPATCH_HOPPER_FNA_BACKWARD_2D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else if constexpr (rank == 3) { \ + DISPATCH_HOPPER_FNA_BACKWARD_3D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Hopper FNA backward kernel dispatch failed! It only supports NA1D, 2D, and 3D!"); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace fna_hopper +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/kernels.h b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/kernels.h new file mode 100644 index 00000000..7008f8d3 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/kernels.h @@ -0,0 +1,7620 @@ +#pragma once + + +#include +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_hopper { + +void hopper_fna1d_backward_float16_64x128x32_Q64_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_backward_float16_128x128x32_Q128_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_backward_float16_64x128x32_Q64_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_backward_float16_128x128x32_Q128_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_backward_float16_64x128x64_Q64_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_backward_float16_128x128x64_Q128_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_backward_float16_64x128x64_Q64_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_backward_float16_128x128x64_Q128_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_backward_float16_64x128x128_Q64_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_backward_float16_64x128x128_Q64_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_backward_bfloat16_64x128x32_Q64_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_backward_bfloat16_128x128x32_Q128_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_backward_bfloat16_64x128x32_Q64_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_backward_bfloat16_128x128x32_Q128_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_backward_bfloat16_64x128x64_Q64_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_backward_bfloat16_128x128x64_Q128_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_backward_bfloat16_64x128x64_Q64_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_backward_bfloat16_128x128x64_Q128_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_backward_bfloat16_64x128x128_Q64_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna1d_backward_bfloat16_64x128x128_Q64_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_64x128x32_Q8x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_64x128x32_Q8x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_128x128x32_Q16x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_128x128x32_Q16x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_64x128x32_Q8x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_64x128x32_Q8x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_128x128x32_Q16x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_128x128x32_Q16x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_64x128x32_Q8x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_64x128x32_Q8x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_128x128x32_Q16x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_128x128x32_Q16x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_64x128x32_Q8x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_64x128x32_Q8x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_128x128x32_Q16x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_128x128x32_Q16x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_64x128x64_Q8x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_64x128x64_Q8x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_128x128x64_Q16x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_128x128x64_Q16x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_64x128x64_Q8x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_64x128x64_Q8x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_128x128x64_Q16x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_128x128x64_Q16x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_64x128x64_Q8x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_64x128x64_Q8x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_128x128x64_Q16x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_128x128x64_Q16x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_64x128x64_Q8x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_64x128x64_Q8x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_128x128x64_Q16x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_128x128x64_Q16x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_64x128x128_Q8x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_64x128x128_Q8x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_64x128x128_Q8x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_64x128x128_Q8x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_64x128x128_Q8x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_64x128x128_Q8x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_64x128x128_Q8x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_float16_64x128x128_Q8x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_64x128x32_Q8x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_64x128x32_Q8x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_128x128x32_Q16x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_128x128x32_Q16x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_64x128x32_Q8x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_64x128x32_Q8x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_128x128x32_Q16x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_128x128x32_Q16x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_64x128x32_Q8x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_64x128x32_Q8x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_128x128x32_Q16x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_128x128x32_Q16x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_64x128x32_Q8x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_64x128x32_Q8x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_128x128x32_Q16x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_128x128x32_Q16x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_64x128x64_Q8x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_64x128x64_Q8x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_128x128x64_Q16x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_128x128x64_Q16x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_64x128x64_Q8x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_64x128x64_Q8x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_128x128x64_Q16x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_128x128x64_Q16x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_64x128x64_Q8x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_64x128x64_Q8x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_128x128x64_Q16x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_128x128x64_Q16x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_64x128x64_Q8x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_64x128x64_Q8x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_128x128x64_Q16x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_128x128x64_Q16x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_64x128x128_Q8x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_64x128x128_Q8x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_64x128x128_Q8x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_64x128x128_Q8x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_64x128x128_Q8x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_64x128x128_Q8x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_64x128x128_Q8x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna2d_backward_bfloat16_64x128x128_Q8x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q2x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q1x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q2x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q1x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q2x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q1x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q2x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q1x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q2x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q1x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q2x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q1x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q2x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q1x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q2x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_float16_64x128x128_Q1x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q2x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q1x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q2x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q1x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q2x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q1x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q2x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q1x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q2x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q1x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q2x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q1x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q2x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q1x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q2x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + +void hopper_fna3d_backward_bfloat16_64x128x128_Q1x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); + + +} // namespace natten +} // namespace cuda +} // namespace fna_hopper +#endif +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/reference/dispatch_cm.h b/natten/csrc/autogen/include/natten_autogen/cuda/reference/dispatch_cm.h new file mode 100644 index 00000000..40175680 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/reference/dispatch_cm.h @@ -0,0 +1,405 @@ +#pragma once + + +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace reference { +#define DISPATCH_REFERENCE_FNA_FORWARD_1D_float32(is_causal, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + natten::cuda::reference::reference_fna1d_float32_causal0(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + natten::cuda::reference::reference_fna1d_float32_causal1(__VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Reference FNA kernel dispatch failed! Causal mask dispatcher got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_REFERENCE_FNA_FORWARD_1D_float16(is_causal, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + natten::cuda::reference::reference_fna1d_float16_causal0(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + natten::cuda::reference::reference_fna1d_float16_causal1(__VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Reference FNA kernel dispatch failed! Causal mask dispatcher got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_REFERENCE_FNA_FORWARD_1D_bfloat16(is_causal, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + natten::cuda::reference::reference_fna1d_bfloat16_causal0(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + natten::cuda::reference::reference_fna1d_bfloat16_causal1(__VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Reference FNA kernel dispatch failed! Causal mask dispatcher got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_REFERENCE_FNA_BACKWARD_1D_float32(is_causal, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + natten::cuda::reference::reference_fna1d_backward_float32_causal0(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + natten::cuda::reference::reference_fna1d_backward_float32_causal1(__VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Reference FNA kernel dispatch failed! Causal mask dispatcher got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_REFERENCE_FNA_BACKWARD_1D_float16(is_causal, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + natten::cuda::reference::reference_fna1d_backward_float16_causal0(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + natten::cuda::reference::reference_fna1d_backward_float16_causal1(__VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Reference FNA kernel dispatch failed! Causal mask dispatcher got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_REFERENCE_FNA_BACKWARD_1D_bfloat16(is_causal, ...) \ + [&] { \ + if (not cute::get<0>(is_causal)) { \ + natten::cuda::reference::reference_fna1d_backward_bfloat16_causal0(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal)) { \ + natten::cuda::reference::reference_fna1d_backward_bfloat16_causal1(__VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Reference FNA kernel dispatch failed! Causal mask dispatcher got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_REFERENCE_FNA_FORWARD_2D_float32(is_causal, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + natten::cuda::reference::reference_fna2d_float32_causal0x0(__VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + natten::cuda::reference::reference_fna2d_float32_causal0x1(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + natten::cuda::reference::reference_fna2d_float32_causal1x0(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + natten::cuda::reference::reference_fna2d_float32_causal1x1(__VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Reference FNA kernel dispatch failed! Causal mask dispatcher got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_REFERENCE_FNA_FORWARD_2D_float16(is_causal, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + natten::cuda::reference::reference_fna2d_float16_causal0x0(__VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + natten::cuda::reference::reference_fna2d_float16_causal0x1(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + natten::cuda::reference::reference_fna2d_float16_causal1x0(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + natten::cuda::reference::reference_fna2d_float16_causal1x1(__VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Reference FNA kernel dispatch failed! Causal mask dispatcher got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_REFERENCE_FNA_FORWARD_2D_bfloat16(is_causal, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + natten::cuda::reference::reference_fna2d_bfloat16_causal0x0(__VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + natten::cuda::reference::reference_fna2d_bfloat16_causal0x1(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + natten::cuda::reference::reference_fna2d_bfloat16_causal1x0(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + natten::cuda::reference::reference_fna2d_bfloat16_causal1x1(__VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Reference FNA kernel dispatch failed! Causal mask dispatcher got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_REFERENCE_FNA_BACKWARD_2D_float32(is_causal, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + natten::cuda::reference::reference_fna2d_backward_float32_causal0x0(__VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + natten::cuda::reference::reference_fna2d_backward_float32_causal0x1(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + natten::cuda::reference::reference_fna2d_backward_float32_causal1x0(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + natten::cuda::reference::reference_fna2d_backward_float32_causal1x1(__VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Reference FNA kernel dispatch failed! Causal mask dispatcher got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_REFERENCE_FNA_BACKWARD_2D_float16(is_causal, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + natten::cuda::reference::reference_fna2d_backward_float16_causal0x0(__VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + natten::cuda::reference::reference_fna2d_backward_float16_causal0x1(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + natten::cuda::reference::reference_fna2d_backward_float16_causal1x0(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + natten::cuda::reference::reference_fna2d_backward_float16_causal1x1(__VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Reference FNA kernel dispatch failed! Causal mask dispatcher got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_REFERENCE_FNA_BACKWARD_2D_bfloat16(is_causal, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + natten::cuda::reference::reference_fna2d_backward_bfloat16_causal0x0(__VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + natten::cuda::reference::reference_fna2d_backward_bfloat16_causal0x1(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal)) { \ + natten::cuda::reference::reference_fna2d_backward_bfloat16_causal1x0(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal)) { \ + natten::cuda::reference::reference_fna2d_backward_bfloat16_causal1x1(__VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Reference FNA kernel dispatch failed! Causal mask dispatcher got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_REFERENCE_FNA_FORWARD_3D_float32(is_causal, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_float32_causal0x0x0(__VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_float32_causal0x0x1(__VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_float32_causal0x1x0(__VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_float32_causal0x1x1(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_float32_causal1x0x0(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_float32_causal1x0x1(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_float32_causal1x1x0(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_float32_causal1x1x1(__VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Reference FNA kernel dispatch failed! Causal mask dispatcher got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_REFERENCE_FNA_FORWARD_3D_float16(is_causal, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_float16_causal0x0x0(__VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_float16_causal0x0x1(__VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_float16_causal0x1x0(__VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_float16_causal0x1x1(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_float16_causal1x0x0(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_float16_causal1x0x1(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_float16_causal1x1x0(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_float16_causal1x1x1(__VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Reference FNA kernel dispatch failed! Causal mask dispatcher got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_REFERENCE_FNA_FORWARD_3D_bfloat16(is_causal, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_bfloat16_causal0x0x0(__VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_bfloat16_causal0x0x1(__VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_bfloat16_causal0x1x0(__VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_bfloat16_causal0x1x1(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_bfloat16_causal1x0x0(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_bfloat16_causal1x0x1(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_bfloat16_causal1x1x0(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_bfloat16_causal1x1x1(__VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Reference FNA kernel dispatch failed! Causal mask dispatcher got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_REFERENCE_FNA_BACKWARD_3D_float32(is_causal, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_backward_float32_causal0x0x0(__VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_backward_float32_causal0x0x1(__VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_backward_float32_causal0x1x0(__VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_backward_float32_causal0x1x1(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_backward_float32_causal1x0x0(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_backward_float32_causal1x0x1(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_backward_float32_causal1x1x0(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_backward_float32_causal1x1x1(__VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Reference FNA kernel dispatch failed! Causal mask dispatcher got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_REFERENCE_FNA_BACKWARD_3D_float16(is_causal, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_backward_float16_causal0x0x0(__VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_backward_float16_causal0x0x1(__VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_backward_float16_causal0x1x0(__VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_backward_float16_causal0x1x1(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_backward_float16_causal1x0x0(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_backward_float16_causal1x0x1(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_backward_float16_causal1x1x0(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_backward_float16_causal1x1x1(__VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Reference FNA kernel dispatch failed! Causal mask dispatcher got invalid causal mask!"); \ + } \ +}(); + +#define DISPATCH_REFERENCE_FNA_BACKWARD_3D_bfloat16(is_causal, ...) \ + [&] { \ + if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_backward_bfloat16_causal0x0x0(__VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_backward_bfloat16_causal0x0x1(__VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_backward_bfloat16_causal0x1x0(__VA_ARGS__); \ + } \ + else if (not cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_backward_bfloat16_causal0x1x1(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_backward_bfloat16_causal1x0x0(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && not cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_backward_bfloat16_causal1x0x1(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && not cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_backward_bfloat16_causal1x1x0(__VA_ARGS__); \ + } \ + else if (cute::get<0>(is_causal) && cute::get<1>(is_causal) && cute::get<2>(is_causal)) { \ + natten::cuda::reference::reference_fna3d_backward_bfloat16_causal1x1x1(__VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Reference FNA kernel dispatch failed! Causal mask dispatcher got invalid causal mask!"); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace reference +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/reference/dispatch_dtype.h b/natten/csrc/autogen/include/natten_autogen/cuda/reference/dispatch_dtype.h new file mode 100644 index 00000000..fe05f1b4 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/reference/dispatch_dtype.h @@ -0,0 +1,123 @@ +#pragma once + + +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace reference { +#define DISPATCH_REFERENCE_FNA_FORWARD_1D(dtype, is_causal, ...) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_REFERENCE_FNA_FORWARD_1D_float32(is_causal, __VA_ARGS__); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_REFERENCE_FNA_FORWARD_1D_float16(is_causal, __VA_ARGS__); \ + } \ + else if (dtype == torch::kBFloat16) { \ + DISPATCH_REFERENCE_FNA_FORWARD_1D_bfloat16(is_causal, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Reference FNA kernel dispatch failed! Reference FNA-1D does not support this data type."); \ + } \ +}(); + +#define DISPATCH_REFERENCE_FNA_BACKWARD_1D(dtype, is_causal, ...) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_REFERENCE_FNA_BACKWARD_1D_float32(is_causal, __VA_ARGS__); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_REFERENCE_FNA_BACKWARD_1D_float16(is_causal, __VA_ARGS__); \ + } \ + else if (dtype == torch::kBFloat16) { \ + DISPATCH_REFERENCE_FNA_BACKWARD_1D_bfloat16(is_causal, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Reference FNA kernel dispatch failed! Reference FNA-1D does not support this data type."); \ + } \ +}(); + +#define DISPATCH_REFERENCE_FNA_FORWARD_2D(dtype, is_causal, ...) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_REFERENCE_FNA_FORWARD_2D_float32(is_causal, __VA_ARGS__); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_REFERENCE_FNA_FORWARD_2D_float16(is_causal, __VA_ARGS__); \ + } \ + else if (dtype == torch::kBFloat16) { \ + DISPATCH_REFERENCE_FNA_FORWARD_2D_bfloat16(is_causal, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Reference FNA kernel dispatch failed! Reference FNA-2D does not support this data type."); \ + } \ +}(); + +#define DISPATCH_REFERENCE_FNA_BACKWARD_2D(dtype, is_causal, ...) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_REFERENCE_FNA_BACKWARD_2D_float32(is_causal, __VA_ARGS__); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_REFERENCE_FNA_BACKWARD_2D_float16(is_causal, __VA_ARGS__); \ + } \ + else if (dtype == torch::kBFloat16) { \ + DISPATCH_REFERENCE_FNA_BACKWARD_2D_bfloat16(is_causal, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Reference FNA kernel dispatch failed! Reference FNA-2D does not support this data type."); \ + } \ +}(); + +#define DISPATCH_REFERENCE_FNA_FORWARD_3D(dtype, is_causal, ...) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_REFERENCE_FNA_FORWARD_3D_float32(is_causal, __VA_ARGS__); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_REFERENCE_FNA_FORWARD_3D_float16(is_causal, __VA_ARGS__); \ + } \ + else if (dtype == torch::kBFloat16) { \ + DISPATCH_REFERENCE_FNA_FORWARD_3D_bfloat16(is_causal, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Reference FNA kernel dispatch failed! Reference FNA-3D does not support this data type."); \ + } \ +}(); + +#define DISPATCH_REFERENCE_FNA_BACKWARD_3D(dtype, is_causal, ...) \ + [&] { \ + if (dtype == torch::kFloat32) { \ + DISPATCH_REFERENCE_FNA_BACKWARD_3D_float32(is_causal, __VA_ARGS__); \ + } \ + else if (dtype == torch::kFloat16) { \ + DISPATCH_REFERENCE_FNA_BACKWARD_3D_float16(is_causal, __VA_ARGS__); \ + } \ + else if (dtype == torch::kBFloat16) { \ + DISPATCH_REFERENCE_FNA_BACKWARD_3D_bfloat16(is_causal, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Reference FNA kernel dispatch failed! Reference FNA-3D does not support this data type."); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace reference +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/reference/interface.h b/natten/csrc/autogen/include/natten_autogen/cuda/reference/interface.h new file mode 100644 index 00000000..376d32a8 --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/reference/interface.h @@ -0,0 +1,59 @@ +#pragma once + + +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace reference { +#define DISPATCH_REFERENCE_FNA_FORWARD(rank, dtype, is_causal, ...) \ + [&] { \ + if constexpr (rank == 1) { \ + DISPATCH_REFERENCE_FNA_FORWARD_1D(dtype, is_causal, __VA_ARGS__); \ + } \ + else if constexpr (rank == 2) { \ + DISPATCH_REFERENCE_FNA_FORWARD_2D(dtype, is_causal, __VA_ARGS__); \ + } \ + else if constexpr (rank == 3) { \ + DISPATCH_REFERENCE_FNA_FORWARD_3D(dtype, is_causal, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Reference FNA kernel dispatch failed! NATTEN only supports NA1D, 2D, and 3D!"); \ + } \ +}(); + +#define DISPATCH_REFERENCE_FNA_BACKWARD(rank, dtype, is_causal, ...) \ + [&] { \ + if constexpr (rank == 1) { \ + DISPATCH_REFERENCE_FNA_BACKWARD_1D(dtype, is_causal, __VA_ARGS__); \ + } \ + else if constexpr (rank == 2) { \ + DISPATCH_REFERENCE_FNA_BACKWARD_2D(dtype, is_causal, __VA_ARGS__); \ + } \ + else if constexpr (rank == 3) { \ + DISPATCH_REFERENCE_FNA_BACKWARD_3D(dtype, is_causal, __VA_ARGS__); \ + } \ + else { \ + throw std::runtime_error("Reference FNA kernel dispatch failed! NATTEN only supports NA1D, 2D, and 3D!"); \ + } \ +}(); + + + +} // namespace natten +} // namespace cuda +} // namespace reference +#endif + diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/reference/kernels.h b/natten/csrc/autogen/include/natten_autogen/cuda/reference/kernels.h new file mode 100644 index 00000000..869b506d --- /dev/null +++ b/natten/csrc/autogen/include/natten_autogen/cuda/reference/kernels.h @@ -0,0 +1,1874 @@ +#pragma once + + +#include +#include +#include +#ifdef NATTEN_WITH_CUTLASS +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace reference { + +void reference_fna1d_float32_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna1d_float32_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna1d_float16_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna1d_float16_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna1d_bfloat16_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna1d_bfloat16_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna1d_backward_float32_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna1d_backward_float32_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna1d_backward_float16_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna1d_backward_float16_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna1d_backward_bfloat16_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna1d_backward_bfloat16_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna2d_float32_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna2d_float32_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna2d_float32_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna2d_float32_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna2d_float16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna2d_float16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna2d_float16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna2d_float16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna2d_bfloat16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna2d_bfloat16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna2d_bfloat16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna2d_bfloat16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna2d_backward_float32_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna2d_backward_float32_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna2d_backward_float32_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna2d_backward_float32_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna2d_backward_float16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna2d_backward_float16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna2d_backward_float16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna2d_backward_float16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna2d_backward_bfloat16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna2d_backward_bfloat16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna2d_backward_bfloat16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna2d_backward_bfloat16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_float32_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_float32_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_float32_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_float32_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_float32_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_float32_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_float32_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_float32_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_float16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_float16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_float16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_float16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_float16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_float16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_float16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_float16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_bfloat16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_bfloat16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_bfloat16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_bfloat16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_bfloat16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_bfloat16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_bfloat16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_bfloat16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_backward_float32_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_backward_float32_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_backward_float32_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_backward_float32_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_backward_float32_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_backward_float32_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_backward_float32_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_backward_float32_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_backward_float16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_backward_float16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_backward_float16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_backward_float16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_backward_float16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_backward_float16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_backward_float16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_backward_float16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_backward_bfloat16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_backward_bfloat16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_backward_bfloat16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_backward_bfloat16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_backward_bfloat16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_backward_bfloat16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_backward_bfloat16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + +void reference_fna3d_backward_bfloat16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream); + + +} // namespace natten +} // namespace cuda +} // namespace reference +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fmha/source_0.cu b/natten/csrc/autogen/src/cuda/blackwell_fmha/source_0.cu new file mode 100644 index 00000000..2e876e5a --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fmha/source_0.cu @@ -0,0 +1,634 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_blackwell { + + + + +void blackwell_fmha_float16_256x128x32( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, GemmShape, false, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, GemmShape, false, cutlass::fmha::collective::ResidualMask, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, GemmShape, false, cutlass::fmha::collective::CausalMask, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, + GemmShape, + /* kIsPersistent= */ false, + cutlass::fmha::collective::ResidualMask, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, + GemmShape, + /* kIsPersistent= */ false, + cutlass::fmha::collective::CausalMask, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void blackwell_fmha_float16_256x128x32_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, GemmShape, true, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, GemmShape, true, cutlass::fmha::collective::ResidualMask, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, GemmShape, true, cutlass::fmha::collective::CausalMask, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, + GemmShape, + /* kIsPersistent= */ true, + cutlass::fmha::collective::ResidualMask, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, + GemmShape, + /* kIsPersistent= */ true, + cutlass::fmha::collective::CausalMask, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void blackwell_fmha_float16_256x128x64( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, GemmShape, false, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, GemmShape, false, cutlass::fmha::collective::ResidualMask, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, GemmShape, false, cutlass::fmha::collective::CausalMask, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, + GemmShape, + /* kIsPersistent= */ false, + cutlass::fmha::collective::ResidualMask, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, + GemmShape, + /* kIsPersistent= */ false, + cutlass::fmha::collective::CausalMask, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void blackwell_fmha_float16_256x128x64_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, GemmShape, true, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, GemmShape, true, cutlass::fmha::collective::ResidualMask, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, GemmShape, true, cutlass::fmha::collective::CausalMask, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, + GemmShape, + /* kIsPersistent= */ true, + cutlass::fmha::collective::ResidualMask, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, + GemmShape, + /* kIsPersistent= */ true, + cutlass::fmha::collective::CausalMask, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void blackwell_fmha_float16_256x128x128( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, GemmShape, false, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, GemmShape, false, cutlass::fmha::collective::ResidualMask, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, GemmShape, false, cutlass::fmha::collective::CausalMask, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, + GemmShape, + /* kIsPersistent= */ false, + cutlass::fmha::collective::ResidualMask, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, + GemmShape, + /* kIsPersistent= */ false, + cutlass::fmha::collective::CausalMask, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void blackwell_fmha_float16_256x128x128_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, GemmShape, true, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, GemmShape, true, cutlass::fmha::collective::ResidualMask, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, GemmShape, true, cutlass::fmha::collective::CausalMask, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, + GemmShape, + /* kIsPersistent= */ true, + cutlass::fmha::collective::ResidualMask, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::half_t, + GemmShape, + /* kIsPersistent= */ true, + cutlass::fmha::collective::CausalMask, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + +} // namespace fmha_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fmha/source_1.cu b/natten/csrc/autogen/src/cuda/blackwell_fmha/source_1.cu new file mode 100644 index 00000000..d8d5a866 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fmha/source_1.cu @@ -0,0 +1,634 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_blackwell { + + + + +void blackwell_fmha_bfloat16_256x128x32( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, GemmShape, false, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, GemmShape, false, cutlass::fmha::collective::ResidualMask, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, GemmShape, false, cutlass::fmha::collective::CausalMask, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, + GemmShape, + /* kIsPersistent= */ false, + cutlass::fmha::collective::ResidualMask, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, + GemmShape, + /* kIsPersistent= */ false, + cutlass::fmha::collective::CausalMask, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void blackwell_fmha_bfloat16_256x128x32_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, GemmShape, true, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, GemmShape, true, cutlass::fmha::collective::ResidualMask, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, GemmShape, true, cutlass::fmha::collective::CausalMask, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, + GemmShape, + /* kIsPersistent= */ true, + cutlass::fmha::collective::ResidualMask, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, + GemmShape, + /* kIsPersistent= */ true, + cutlass::fmha::collective::CausalMask, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void blackwell_fmha_bfloat16_256x128x64( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, GemmShape, false, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, GemmShape, false, cutlass::fmha::collective::ResidualMask, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, GemmShape, false, cutlass::fmha::collective::CausalMask, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, + GemmShape, + /* kIsPersistent= */ false, + cutlass::fmha::collective::ResidualMask, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, + GemmShape, + /* kIsPersistent= */ false, + cutlass::fmha::collective::CausalMask, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void blackwell_fmha_bfloat16_256x128x64_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, GemmShape, true, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, GemmShape, true, cutlass::fmha::collective::ResidualMask, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, GemmShape, true, cutlass::fmha::collective::CausalMask, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, + GemmShape, + /* kIsPersistent= */ true, + cutlass::fmha::collective::ResidualMask, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, + GemmShape, + /* kIsPersistent= */ true, + cutlass::fmha::collective::CausalMask, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void blackwell_fmha_bfloat16_256x128x128( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, GemmShape, false, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, GemmShape, false, cutlass::fmha::collective::ResidualMask, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, GemmShape, false, cutlass::fmha::collective::CausalMask, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, + GemmShape, + /* kIsPersistent= */ false, + cutlass::fmha::collective::ResidualMask, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, + GemmShape, + /* kIsPersistent= */ false, + cutlass::fmha::collective::CausalMask, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void blackwell_fmha_bfloat16_256x128x128_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, GemmShape, true, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, GemmShape, true, cutlass::fmha::collective::ResidualMask, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, GemmShape, true, cutlass::fmha::collective::CausalMask, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, + GemmShape, + /* kIsPersistent= */ true, + cutlass::fmha::collective::ResidualMask, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::bfloat16_t, + GemmShape, + /* kIsPersistent= */ true, + cutlass::fmha::collective::CausalMask, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + +} // namespace fmha_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fmha/source_2.cu b/natten/csrc/autogen/src/cuda/blackwell_fmha/source_2.cu new file mode 100644 index 00000000..9fe09516 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fmha/source_2.cu @@ -0,0 +1,634 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_blackwell { + + + + +void blackwell_fmha_e4m3_256x128x32( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, GemmShape, false, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, GemmShape, false, cutlass::fmha::collective::ResidualMask, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, GemmShape, false, cutlass::fmha::collective::CausalMask, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, + GemmShape, + /* kIsPersistent= */ false, + cutlass::fmha::collective::ResidualMask, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, + GemmShape, + /* kIsPersistent= */ false, + cutlass::fmha::collective::CausalMask, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void blackwell_fmha_e4m3_256x128x32_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, GemmShape, true, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, GemmShape, true, cutlass::fmha::collective::ResidualMask, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, GemmShape, true, cutlass::fmha::collective::CausalMask, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, + GemmShape, + /* kIsPersistent= */ true, + cutlass::fmha::collective::ResidualMask, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, + GemmShape, + /* kIsPersistent= */ true, + cutlass::fmha::collective::CausalMask, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void blackwell_fmha_e4m3_256x128x64( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, GemmShape, false, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, GemmShape, false, cutlass::fmha::collective::ResidualMask, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, GemmShape, false, cutlass::fmha::collective::CausalMask, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, + GemmShape, + /* kIsPersistent= */ false, + cutlass::fmha::collective::ResidualMask, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, + GemmShape, + /* kIsPersistent= */ false, + cutlass::fmha::collective::CausalMask, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void blackwell_fmha_e4m3_256x128x64_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, GemmShape, true, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, GemmShape, true, cutlass::fmha::collective::ResidualMask, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, GemmShape, true, cutlass::fmha::collective::CausalMask, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, + GemmShape, + /* kIsPersistent= */ true, + cutlass::fmha::collective::ResidualMask, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, + GemmShape, + /* kIsPersistent= */ true, + cutlass::fmha::collective::CausalMask, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void blackwell_fmha_e4m3_256x128x128( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, GemmShape, false, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, GemmShape, false, cutlass::fmha::collective::ResidualMask, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, GemmShape, false, cutlass::fmha::collective::CausalMask, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, + GemmShape, + /* kIsPersistent= */ false, + cutlass::fmha::collective::ResidualMask, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, + GemmShape, + /* kIsPersistent= */ false, + cutlass::fmha::collective::CausalMask, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void blackwell_fmha_e4m3_256x128x128_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, GemmShape, true, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, GemmShape, true, cutlass::fmha::collective::ResidualMask, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, GemmShape, true, cutlass::fmha::collective::CausalMask, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, + GemmShape, + /* kIsPersistent= */ true, + cutlass::fmha::collective::ResidualMask, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e4m3_t, + GemmShape, + /* kIsPersistent= */ true, + cutlass::fmha::collective::CausalMask, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + +} // namespace fmha_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fmha/source_3.cu b/natten/csrc/autogen/src/cuda/blackwell_fmha/source_3.cu new file mode 100644 index 00000000..dbdafcbd --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fmha/source_3.cu @@ -0,0 +1,634 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_blackwell { + + + + +void blackwell_fmha_e5m2_256x128x32( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, GemmShape, false, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, GemmShape, false, cutlass::fmha::collective::ResidualMask, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, GemmShape, false, cutlass::fmha::collective::CausalMask, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, + GemmShape, + /* kIsPersistent= */ false, + cutlass::fmha::collective::ResidualMask, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, + GemmShape, + /* kIsPersistent= */ false, + cutlass::fmha::collective::CausalMask, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void blackwell_fmha_e5m2_256x128x32_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, GemmShape, true, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, GemmShape, true, cutlass::fmha::collective::ResidualMask, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, GemmShape, true, cutlass::fmha::collective::CausalMask, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, + GemmShape, + /* kIsPersistent= */ true, + cutlass::fmha::collective::ResidualMask, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, + GemmShape, + /* kIsPersistent= */ true, + cutlass::fmha::collective::CausalMask, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void blackwell_fmha_e5m2_256x128x64( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, GemmShape, false, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, GemmShape, false, cutlass::fmha::collective::ResidualMask, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, GemmShape, false, cutlass::fmha::collective::CausalMask, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, + GemmShape, + /* kIsPersistent= */ false, + cutlass::fmha::collective::ResidualMask, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, + GemmShape, + /* kIsPersistent= */ false, + cutlass::fmha::collective::CausalMask, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void blackwell_fmha_e5m2_256x128x64_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, GemmShape, true, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, GemmShape, true, cutlass::fmha::collective::ResidualMask, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, GemmShape, true, cutlass::fmha::collective::CausalMask, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, + GemmShape, + /* kIsPersistent= */ true, + cutlass::fmha::collective::ResidualMask, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, + GemmShape, + /* kIsPersistent= */ true, + cutlass::fmha::collective::CausalMask, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void blackwell_fmha_e5m2_256x128x128( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, GemmShape, false, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, GemmShape, false, cutlass::fmha::collective::ResidualMask, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, GemmShape, false, cutlass::fmha::collective::CausalMask, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, + GemmShape, + /* kIsPersistent= */ false, + cutlass::fmha::collective::ResidualMask, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, + GemmShape, + /* kIsPersistent= */ false, + cutlass::fmha::collective::CausalMask, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void blackwell_fmha_e5m2_256x128x128_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, GemmShape, true, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, GemmShape, true, cutlass::fmha::collective::ResidualMask, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, GemmShape, true, cutlass::fmha::collective::CausalMask, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, + GemmShape, + /* kIsPersistent= */ true, + cutlass::fmha::collective::ResidualMask, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + cutlass::float_e5m2_t, + GemmShape, + /* kIsPersistent= */ true, + cutlass::fmha::collective::CausalMask, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + +} // namespace fmha_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_0.cu b/natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_0.cu new file mode 100644 index 00000000..1e08b3cc --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_0.cu @@ -0,0 +1,288 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_blackwell { + + + + +void blackwell_fmha_backward_float16_128x128x32( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + bool deterministic, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + auto run = [&](auto kIsDeterministic) { + constexpr bool IsDeterministic = decltype(kIsDeterministic)::value; + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::half_t, GemmShape, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false, IsDeterministic>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::half_t, GemmShape, cutlass::fmha::collective::ResidualMaskForBackward, /* kIsVarlen= */ false, IsDeterministic>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::half_t, GemmShape, cutlass::fmha::collective::CausalForBackwardMask, /* kIsVarlen= */ false, IsDeterministic>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::half_t, + GemmShape, + cutlass::fmha::collective::ResidualMaskForBackward, + /* kIsVarlen= */ true, IsDeterministic>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::half_t, + GemmShape, + cutlass::fmha::collective::CausalForBackwardMask, + /* kIsVarlen= */ true, IsDeterministic>; + + int* dq_semaphore_ptr = nullptr; + at::Tensor dq_semaphore; + if constexpr (IsDeterministic) { + auto kBlockM = cute::get<0>(GemmShape{}); + int effective_seqlen_q = is_varlen ? max_seqlen_Q : seqlen_q; + dq_semaphore = torch::zeros( + {(effective_seqlen_q + kBlockM - 1) / kBlockM, batch_size, heads_q}, + tensor_options.dtype(at::ScalarType::Int)); + dq_semaphore_ptr = static_cast(dq_semaphore.data_ptr()); + } + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + dq_semaphore_ptr, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && + seqlen_q % get<0>(GemmShape{}) == 0 && seqlen_k % get<1>(GemmShape{}) == 0; + + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } + }; + + if (deterministic) { + run(std::true_type{}); + } else { + run(std::false_type{}); + } +} + + + + + +void blackwell_fmha_backward_float16_128x128x64( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + bool deterministic, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + auto run = [&](auto kIsDeterministic) { + constexpr bool IsDeterministic = decltype(kIsDeterministic)::value; + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::half_t, GemmShape, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false, IsDeterministic>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::half_t, GemmShape, cutlass::fmha::collective::ResidualMaskForBackward, /* kIsVarlen= */ false, IsDeterministic>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::half_t, GemmShape, cutlass::fmha::collective::CausalForBackwardMask, /* kIsVarlen= */ false, IsDeterministic>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::half_t, + GemmShape, + cutlass::fmha::collective::ResidualMaskForBackward, + /* kIsVarlen= */ true, IsDeterministic>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::half_t, + GemmShape, + cutlass::fmha::collective::CausalForBackwardMask, + /* kIsVarlen= */ true, IsDeterministic>; + + int* dq_semaphore_ptr = nullptr; + at::Tensor dq_semaphore; + if constexpr (IsDeterministic) { + auto kBlockM = cute::get<0>(GemmShape{}); + int effective_seqlen_q = is_varlen ? max_seqlen_Q : seqlen_q; + dq_semaphore = torch::zeros( + {(effective_seqlen_q + kBlockM - 1) / kBlockM, batch_size, heads_q}, + tensor_options.dtype(at::ScalarType::Int)); + dq_semaphore_ptr = static_cast(dq_semaphore.data_ptr()); + } + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + dq_semaphore_ptr, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && + seqlen_q % get<0>(GemmShape{}) == 0 && seqlen_k % get<1>(GemmShape{}) == 0; + + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } + }; + + if (deterministic) { + run(std::true_type{}); + } else { + run(std::false_type{}); + } +} + + +} // namespace fmha_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_1.cu b/natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_1.cu new file mode 100644 index 00000000..9cb855d9 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_1.cu @@ -0,0 +1,288 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_blackwell { + + + + +void blackwell_fmha_backward_float16_128x128x128( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + bool deterministic, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + auto run = [&](auto kIsDeterministic) { + constexpr bool IsDeterministic = decltype(kIsDeterministic)::value; + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::half_t, GemmShape, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false, IsDeterministic>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::half_t, GemmShape, cutlass::fmha::collective::ResidualMaskForBackward, /* kIsVarlen= */ false, IsDeterministic>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::half_t, GemmShape, cutlass::fmha::collective::CausalForBackwardMask, /* kIsVarlen= */ false, IsDeterministic>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::half_t, + GemmShape, + cutlass::fmha::collective::ResidualMaskForBackward, + /* kIsVarlen= */ true, IsDeterministic>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::half_t, + GemmShape, + cutlass::fmha::collective::CausalForBackwardMask, + /* kIsVarlen= */ true, IsDeterministic>; + + int* dq_semaphore_ptr = nullptr; + at::Tensor dq_semaphore; + if constexpr (IsDeterministic) { + auto kBlockM = cute::get<0>(GemmShape{}); + int effective_seqlen_q = is_varlen ? max_seqlen_Q : seqlen_q; + dq_semaphore = torch::zeros( + {(effective_seqlen_q + kBlockM - 1) / kBlockM, batch_size, heads_q}, + tensor_options.dtype(at::ScalarType::Int)); + dq_semaphore_ptr = static_cast(dq_semaphore.data_ptr()); + } + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + dq_semaphore_ptr, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && + seqlen_q % get<0>(GemmShape{}) == 0 && seqlen_k % get<1>(GemmShape{}) == 0; + + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } + }; + + if (deterministic) { + run(std::true_type{}); + } else { + run(std::false_type{}); + } +} + + + + + +void blackwell_fmha_backward_bfloat16_128x128x32( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + bool deterministic, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + auto run = [&](auto kIsDeterministic) { + constexpr bool IsDeterministic = decltype(kIsDeterministic)::value; + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::bfloat16_t, GemmShape, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false, IsDeterministic>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::bfloat16_t, GemmShape, cutlass::fmha::collective::ResidualMaskForBackward, /* kIsVarlen= */ false, IsDeterministic>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::bfloat16_t, GemmShape, cutlass::fmha::collective::CausalForBackwardMask, /* kIsVarlen= */ false, IsDeterministic>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::bfloat16_t, + GemmShape, + cutlass::fmha::collective::ResidualMaskForBackward, + /* kIsVarlen= */ true, IsDeterministic>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::bfloat16_t, + GemmShape, + cutlass::fmha::collective::CausalForBackwardMask, + /* kIsVarlen= */ true, IsDeterministic>; + + int* dq_semaphore_ptr = nullptr; + at::Tensor dq_semaphore; + if constexpr (IsDeterministic) { + auto kBlockM = cute::get<0>(GemmShape{}); + int effective_seqlen_q = is_varlen ? max_seqlen_Q : seqlen_q; + dq_semaphore = torch::zeros( + {(effective_seqlen_q + kBlockM - 1) / kBlockM, batch_size, heads_q}, + tensor_options.dtype(at::ScalarType::Int)); + dq_semaphore_ptr = static_cast(dq_semaphore.data_ptr()); + } + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + dq_semaphore_ptr, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && + seqlen_q % get<0>(GemmShape{}) == 0 && seqlen_k % get<1>(GemmShape{}) == 0; + + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } + }; + + if (deterministic) { + run(std::true_type{}); + } else { + run(std::false_type{}); + } +} + + +} // namespace fmha_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_2.cu b/natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_2.cu new file mode 100644 index 00000000..0ebf4405 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_2.cu @@ -0,0 +1,155 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_blackwell { + + + + +void blackwell_fmha_backward_bfloat16_128x128x64( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + bool deterministic, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + auto run = [&](auto kIsDeterministic) { + constexpr bool IsDeterministic = decltype(kIsDeterministic)::value; + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::bfloat16_t, GemmShape, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false, IsDeterministic>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::bfloat16_t, GemmShape, cutlass::fmha::collective::ResidualMaskForBackward, /* kIsVarlen= */ false, IsDeterministic>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::bfloat16_t, GemmShape, cutlass::fmha::collective::CausalForBackwardMask, /* kIsVarlen= */ false, IsDeterministic>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::bfloat16_t, + GemmShape, + cutlass::fmha::collective::ResidualMaskForBackward, + /* kIsVarlen= */ true, IsDeterministic>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::bfloat16_t, + GemmShape, + cutlass::fmha::collective::CausalForBackwardMask, + /* kIsVarlen= */ true, IsDeterministic>; + + int* dq_semaphore_ptr = nullptr; + at::Tensor dq_semaphore; + if constexpr (IsDeterministic) { + auto kBlockM = cute::get<0>(GemmShape{}); + int effective_seqlen_q = is_varlen ? max_seqlen_Q : seqlen_q; + dq_semaphore = torch::zeros( + {(effective_seqlen_q + kBlockM - 1) / kBlockM, batch_size, heads_q}, + tensor_options.dtype(at::ScalarType::Int)); + dq_semaphore_ptr = static_cast(dq_semaphore.data_ptr()); + } + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + dq_semaphore_ptr, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && + seqlen_q % get<0>(GemmShape{}) == 0 && seqlen_k % get<1>(GemmShape{}) == 0; + + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } + }; + + if (deterministic) { + run(std::true_type{}); + } else { + run(std::false_type{}); + } +} + + +} // namespace fmha_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_3.cu b/natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_3.cu new file mode 100644 index 00000000..b934754f --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_3.cu @@ -0,0 +1,155 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_blackwell { + + + + +void blackwell_fmha_backward_bfloat16_128x128x128( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + bool deterministic, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + auto run = [&](auto kIsDeterministic) { + constexpr bool IsDeterministic = decltype(kIsDeterministic)::value; + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::bfloat16_t, GemmShape, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false, IsDeterministic>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::bfloat16_t, GemmShape, cutlass::fmha::collective::ResidualMaskForBackward, /* kIsVarlen= */ false, IsDeterministic>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::bfloat16_t, GemmShape, cutlass::fmha::collective::CausalForBackwardMask, /* kIsVarlen= */ false, IsDeterministic>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::bfloat16_t, + GemmShape, + cutlass::fmha::collective::ResidualMaskForBackward, + /* kIsVarlen= */ true, IsDeterministic>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelBackward< + cutlass::bfloat16_t, + GemmShape, + cutlass::fmha::collective::CausalForBackwardMask, + /* kIsVarlen= */ true, IsDeterministic>; + + int* dq_semaphore_ptr = nullptr; + at::Tensor dq_semaphore; + if constexpr (IsDeterministic) { + auto kBlockM = cute::get<0>(GemmShape{}); + int effective_seqlen_q = is_varlen ? max_seqlen_Q : seqlen_q; + dq_semaphore = torch::zeros( + {(effective_seqlen_q + kBlockM - 1) / kBlockM, batch_size, heads_q}, + tensor_options.dtype(at::ScalarType::Int)); + dq_semaphore_ptr = static_cast(dq_semaphore.data_ptr()); + } + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + dq_semaphore_ptr, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && + seqlen_q % get<0>(GemmShape{}) == 0 && seqlen_k % get<1>(GemmShape{}) == 0; + + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } + }; + + if (deterministic) { + run(std::true_type{}); + } else { + run(std::false_type{}); + } +} + + +} // namespace fmha_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna/source_0.cu b/natten/csrc/autogen/src/cuda/blackwell_fna/source_0.cu new file mode 100644 index 00000000..92f2edea --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna/source_0.cu @@ -0,0 +1,4118 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna1d_float16_256x128x32_Q256_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_float16_256x128x32_Q256_KV128_causal0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_float16_256x128x32_Q256_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_float16_256x128x32_Q256_KV128_causal1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_float16_256x128x64_Q256_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_float16_256x128x64_Q256_KV128_causal0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_float16_256x128x64_Q256_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_float16_256x128x64_Q256_KV128_causal1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_float16_256x128x128_Q256_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_float16_256x128x128_Q256_KV128_causal0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_float16_256x128x128_Q256_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_float16_256x128x128_Q256_KV128_causal1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_bfloat16_256x128x32_Q256_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_bfloat16_256x128x32_Q256_KV128_causal0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_bfloat16_256x128x32_Q256_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_bfloat16_256x128x32_Q256_KV128_causal1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_bfloat16_256x128x64_Q256_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_bfloat16_256x128x64_Q256_KV128_causal0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_bfloat16_256x128x64_Q256_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_bfloat16_256x128x64_Q256_KV128_causal1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_bfloat16_256x128x128_Q256_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_bfloat16_256x128x128_Q256_KV128_causal0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_bfloat16_256x128x128_Q256_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_bfloat16_256x128x128_Q256_KV128_causal1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_e4m3_256x128x32_Q256_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_e4m3_256x128x32_Q256_KV128_causal0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_e4m3_256x128x32_Q256_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_e4m3_256x128x32_Q256_KV128_causal1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_e4m3_256x128x64_Q256_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_e4m3_256x128x64_Q256_KV128_causal0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_e4m3_256x128x64_Q256_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_e4m3_256x128x64_Q256_KV128_causal1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_e4m3_256x128x128_Q256_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_e4m3_256x128x128_Q256_KV128_causal0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_e4m3_256x128x128_Q256_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_e4m3_256x128x128_Q256_KV128_causal1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_e5m2_256x128x32_Q256_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_e5m2_256x128x32_Q256_KV128_causal0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_e5m2_256x128x32_Q256_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_e5m2_256x128x32_Q256_KV128_causal1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_e5m2_256x128x64_Q256_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_e5m2_256x128x64_Q256_KV128_causal0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_e5m2_256x128x64_Q256_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_e5m2_256x128x64_Q256_KV128_causal1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_e5m2_256x128x128_Q256_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_e5m2_256x128x128_Q256_KV128_causal0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_e5m2_256x128x128_Q256_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_e5m2_256x128x128_Q256_KV128_causal1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV16x8_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV4x32_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV4x32_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV16x8_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV4x32_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV4x32_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna/source_1.cu b/natten/csrc/autogen/src/cuda/blackwell_fna/source_1.cu new file mode 100644 index 00000000..a7b9c4e5 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna/source_1.cu @@ -0,0 +1,4118 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV16x8_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV4x32_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV4x32_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV16x8_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q16x16_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV4x32_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x32_Q8x32_KV4x32_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV16x8_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV4x32_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV4x32_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV16x8_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV4x32_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV4x32_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV16x8_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV4x32_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV4x32_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV16x8_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q16x16_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV4x32_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x64_Q8x32_KV4x32_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV16x8_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV4x32_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV4x32_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV16x8_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV4x32_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV4x32_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna/source_10.cu b/natten/csrc/autogen/src/cuda/blackwell_fna/source_10.cu new file mode 100644 index 00000000..5a306d82 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna/source_10.cu @@ -0,0 +1,4118 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna/source_11.cu b/natten/csrc/autogen/src/cuda/blackwell_fna/source_11.cu new file mode 100644 index 00000000..56108d25 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna/source_11.cu @@ -0,0 +1,4118 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q8x4x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x8x16_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x4x16_KV2x4x16_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q2x16x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x128_Q4x8x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna/source_12.cu b/natten/csrc/autogen/src/cuda/blackwell_fna/source_12.cu new file mode 100644 index 00000000..22ba0c81 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna/source_12.cu @@ -0,0 +1,4054 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna/source_13.cu b/natten/csrc/autogen/src/cuda/blackwell_fna/source_13.cu new file mode 100644 index 00000000..b36cda5b --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna/source_13.cu @@ -0,0 +1,4054 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q8x4x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x8x16_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x4x16_KV2x4x16_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q2x16x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x32_Q4x8x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna/source_14.cu b/natten/csrc/autogen/src/cuda/blackwell_fna/source_14.cu new file mode 100644 index 00000000..984beb7a --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna/source_14.cu @@ -0,0 +1,4054 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna/source_15.cu b/natten/csrc/autogen/src/cuda/blackwell_fna/source_15.cu new file mode 100644 index 00000000..67688b55 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna/source_15.cu @@ -0,0 +1,4054 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q8x4x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x8x16_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x4x16_KV2x4x16_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q2x16x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x64_Q4x8x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna/source_16.cu b/natten/csrc/autogen/src/cuda/blackwell_fna/source_16.cu new file mode 100644 index 00000000..3f1975b5 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna/source_16.cu @@ -0,0 +1,4054 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna/source_17.cu b/natten/csrc/autogen/src/cuda/blackwell_fna/source_17.cu new file mode 100644 index 00000000..a38a49cf --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna/source_17.cu @@ -0,0 +1,4054 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q8x4x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x8x16_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x4x16_KV2x4x16_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q2x16x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_bfloat16_256x128x128_Q4x8x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna/source_18.cu b/natten/csrc/autogen/src/cuda/blackwell_fna/source_18.cu new file mode 100644 index 00000000..1b6035ae --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna/source_18.cu @@ -0,0 +1,4054 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q8x4x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna/source_19.cu b/natten/csrc/autogen/src/cuda/blackwell_fna/source_19.cu new file mode 100644 index 00000000..d7339d89 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna/source_19.cu @@ -0,0 +1,4054 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x8x16_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x4x16_KV2x4x16_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q2x16x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x32_Q4x8x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna/source_2.cu b/natten/csrc/autogen/src/cuda/blackwell_fna/source_2.cu new file mode 100644 index 00000000..5a8803d9 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna/source_2.cu @@ -0,0 +1,4118 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV16x8_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV4x32_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV4x32_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV16x8_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q16x16_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV4x32_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_float16_256x128x128_Q8x32_KV4x32_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV16x8_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV4x32_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV4x32_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV16x8_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV4x32_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV4x32_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV16x8_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV4x32_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV4x32_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV16x8_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q16x16_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV4x32_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x32_Q8x32_KV4x32_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV16x8_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV4x32_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV4x32_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV16x8_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV4x32_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV4x32_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna/source_20.cu b/natten/csrc/autogen/src/cuda/blackwell_fna/source_20.cu new file mode 100644 index 00000000..88715cb6 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna/source_20.cu @@ -0,0 +1,4054 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q8x4x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x8x16_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x4x16_KV2x4x16_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q2x16x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x64_Q4x8x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna/source_21.cu b/natten/csrc/autogen/src/cuda/blackwell_fna/source_21.cu new file mode 100644 index 00000000..688c0c1c --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna/source_21.cu @@ -0,0 +1,4054 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna/source_22.cu b/natten/csrc/autogen/src/cuda/blackwell_fna/source_22.cu new file mode 100644 index 00000000..d556adaa --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna/source_22.cu @@ -0,0 +1,4054 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q8x4x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x8x16_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x4x16_KV2x4x16_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q2x16x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e4m3_256x128x128_Q4x8x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna/source_23.cu b/natten/csrc/autogen/src/cuda/blackwell_fna/source_23.cu new file mode 100644 index 00000000..eeb23d49 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna/source_23.cu @@ -0,0 +1,4054 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna/source_24.cu b/natten/csrc/autogen/src/cuda/blackwell_fna/source_24.cu new file mode 100644 index 00000000..b13a22d5 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna/source_24.cu @@ -0,0 +1,4054 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q8x4x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x8x16_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x4x16_KV2x4x16_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q2x16x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x32_Q4x8x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna/source_25.cu b/natten/csrc/autogen/src/cuda/blackwell_fna/source_25.cu new file mode 100644 index 00000000..e4e228fb --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna/source_25.cu @@ -0,0 +1,4054 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna/source_26.cu b/natten/csrc/autogen/src/cuda/blackwell_fna/source_26.cu new file mode 100644 index 00000000..ecd1705a --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna/source_26.cu @@ -0,0 +1,4054 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q8x4x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x8x16_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x4x16_KV2x4x16_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q2x16x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x64_Q4x8x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna/source_27.cu b/natten/csrc/autogen/src/cuda/blackwell_fna/source_27.cu new file mode 100644 index 00000000..73fea315 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna/source_27.cu @@ -0,0 +1,4054 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q8x4x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x8x16_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x4x16_KV2x4x16_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q2x16x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_e5m2_256x128x128_Q4x8x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna/source_3.cu b/natten/csrc/autogen/src/cuda/blackwell_fna/source_3.cu new file mode 100644 index 00000000..d0285ec4 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna/source_3.cu @@ -0,0 +1,4118 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV16x8_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV4x32_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV4x32_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV16x8_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q16x16_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV4x32_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x64_Q8x32_KV4x32_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV16x8_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV4x32_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV4x32_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV16x8_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV4x32_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV4x32_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV16x8_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV4x32_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV4x32_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV16x8_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q16x16_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV4x32_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_bfloat16_256x128x128_Q8x32_KV4x32_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV16x8_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV4x32_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV4x32_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV16x8_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV4x32_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV4x32_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna/source_4.cu b/natten/csrc/autogen/src/cuda/blackwell_fna/source_4.cu new file mode 100644 index 00000000..ada00c77 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna/source_4.cu @@ -0,0 +1,4118 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV16x8_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV4x32_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV4x32_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV16x8_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q16x16_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV4x32_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x32_Q8x32_KV4x32_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV16x8_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV4x32_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV4x32_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV16x8_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV4x32_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV4x32_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV16x8_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV4x32_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV4x32_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV16x8_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q16x16_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV4x32_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x64_Q8x32_KV4x32_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV16x8_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV4x32_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV4x32_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV16x8_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV4x32_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV4x32_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna/source_5.cu b/natten/csrc/autogen/src/cuda/blackwell_fna/source_5.cu new file mode 100644 index 00000000..fda6da44 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna/source_5.cu @@ -0,0 +1,4118 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV16x8_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV4x32_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV4x32_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV16x8_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q16x16_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV4x32_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e4m3_256x128x128_Q8x32_KV4x32_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e4m3_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV16x8_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV4x32_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV4x32_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV16x8_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV4x32_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV4x32_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV16x8_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV4x32_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV4x32_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV16x8_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q16x16_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV4x32_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x32_Q8x32_KV4x32_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV16x8_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV4x32_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV4x32_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV16x8_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV4x32_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV4x32_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna/source_6.cu b/natten/csrc/autogen/src/cuda/blackwell_fna/source_6.cu new file mode 100644 index 00000000..d86bbfdb --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna/source_6.cu @@ -0,0 +1,4118 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV16x8_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV4x32_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV4x32_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV16x8_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q16x16_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV4x32_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x64_Q8x32_KV4x32_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV16x8_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV8x16_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV4x32_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV4x32_causal0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV16x8_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV8x16_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV4x32_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV4x32_causal0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV16x8_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV8x16_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV4x32_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV4x32_causal1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV16x8_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q16x16_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV8x16_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV4x32_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_e5m2_256x128x128_Q8x32_KV4x32_causal1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<32>>; + using KVTileShape = cute::tuple, cute::Int<32>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::float_e5m2_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna/source_7.cu b/natten/csrc/autogen/src/cuda/blackwell_fna/source_7.cu new file mode 100644 index 00000000..e05855db --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna/source_7.cu @@ -0,0 +1,4118 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna/source_8.cu b/natten/csrc/autogen/src/cuda/blackwell_fna/source_8.cu new file mode 100644 index 00000000..73d260b8 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna/source_8.cu @@ -0,0 +1,4118 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q8x4x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV4x4x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x8x16_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x4x16_KV2x4x16_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q2x16x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x32_Q4x8x8_KV2x8x8_causal1x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal0x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal0x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna/source_9.cu b/natten/csrc/autogen/src/cuda/blackwell_fna/source_9.cu new file mode 100644 index 00000000..7e37b4f1 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna/source_9.cu @@ -0,0 +1,4118 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal0x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal0x1x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal1x0x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x8x8_KV2x8x8_causal1x0x1_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q8x4x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV4x4x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x8x16_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q4x4x16_KV2x4x16_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, false>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_float16_256x128x64_Q2x16x8_KV2x8x8_causal1x1x0_persistent( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config, true>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_0.cu b/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_0.cu new file mode 100644 index 00000000..c17cbc4f --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_0.cu @@ -0,0 +1,2830 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna1d_backward_float16_128x128x32_Q128_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_backward_float16_128x128x32_Q128_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_backward_float16_128x128x64_Q128_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_backward_float16_128x128x64_Q128_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_backward_float16_128x128x128_Q128_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_backward_float16_128x128x128_Q128_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_backward_bfloat16_128x128x32_Q128_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_backward_bfloat16_128x128x32_Q128_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_backward_bfloat16_128x128x64_Q128_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_backward_bfloat16_128x128x64_Q128_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_backward_bfloat16_128x128x128_Q128_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna1d_backward_bfloat16_128x128x128_Q128_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x32_Q16x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x32_Q16x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x32_Q8x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x32_Q8x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x32_Q16x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x32_Q16x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x32_Q8x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x32_Q8x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x32_Q16x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x32_Q16x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x32_Q8x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x32_Q8x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x32_Q16x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x32_Q16x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x32_Q8x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x32_Q8x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x64_Q16x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x64_Q16x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x64_Q8x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x64_Q8x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x64_Q16x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x64_Q16x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x64_Q8x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x64_Q8x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x64_Q16x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x64_Q16x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x64_Q8x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_1.cu b/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_1.cu new file mode 100644 index 00000000..ee2d9caa --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_1.cu @@ -0,0 +1,2830 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna2d_backward_float16_128x128x64_Q8x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x64_Q16x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x64_Q16x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x64_Q8x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x64_Q8x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x128_Q16x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x128_Q16x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x128_Q8x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x128_Q8x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x128_Q16x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x128_Q16x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x128_Q8x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x128_Q8x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x128_Q16x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x128_Q16x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x128_Q8x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x128_Q8x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x128_Q16x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x128_Q16x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x128_Q8x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_float16_128x128x128_Q8x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q16x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q16x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q8x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q8x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q16x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q16x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q8x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q8x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q16x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q16x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q8x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q8x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q16x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q16x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q8x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x32_Q8x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q16x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q16x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_10.cu b/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_10.cu new file mode 100644 index 00000000..a27111ea --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_10.cu @@ -0,0 +1,2758 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x2x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x2x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x2x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x2x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_11.cu b/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_11.cu new file mode 100644 index 00000000..79e86328 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_11.cu @@ -0,0 +1,2758 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x2x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x2x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x2x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x2x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q1x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_12.cu b/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_12.cu new file mode 100644 index 00000000..33407bd4 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_12.cu @@ -0,0 +1,2758 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x2x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x64_Q2x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x2x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x2x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x2x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x2x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_13.cu b/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_13.cu new file mode 100644 index 00000000..2813c6e2 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_13.cu @@ -0,0 +1,2758 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x2x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x2x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x2x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q1x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x2x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q4x4x8_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x128_Q2x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_2.cu b/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_2.cu new file mode 100644 index 00000000..a96e52c2 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_2.cu @@ -0,0 +1,2830 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q8x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q8x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q16x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q16x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q8x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q8x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q16x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q16x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q8x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q8x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q16x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q16x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q8x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x64_Q8x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q16x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q16x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q8x16_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q8x16_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q16x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q16x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q8x16_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q8x16_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q16x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q16x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q8x16_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q8x16_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q16x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q16x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q8x16_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna2d_backward_bfloat16_128x128x128_Q8x16_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q2x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x2x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_3.cu b/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_3.cu new file mode 100644 index 00000000..1ffc9693 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_3.cu @@ -0,0 +1,2830 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q2x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x2x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q2x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x2x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q2x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x2x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q2x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x2x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_4.cu b/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_4.cu new file mode 100644 index 00000000..6658abd0 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_4.cu @@ -0,0 +1,2830 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q2x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x2x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q2x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x2x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q1x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q2x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x2x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x32_Q2x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q2x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x2x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q2x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_5.cu b/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_5.cu new file mode 100644 index 00000000..a32fe9a2 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_5.cu @@ -0,0 +1,2830 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x2x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q2x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x2x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q2x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x2x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q2x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x2x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q2x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x2x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_6.cu b/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_6.cu new file mode 100644 index 00000000..e8c36cec --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_6.cu @@ -0,0 +1,2830 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q2x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x2x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q1x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q2x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x2x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x64_Q2x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q2x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x2x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q2x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x2x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_7.cu b/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_7.cu new file mode 100644 index 00000000..4553138b --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_7.cu @@ -0,0 +1,2830 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q2x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x2x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q2x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x2x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q2x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x2x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q2x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x2x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q2x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_8.cu b/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_8.cu new file mode 100644 index 00000000..7e5b7869 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_8.cu @@ -0,0 +1,2758 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x2x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q1x8x16_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q2x4x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x2x16_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q4x4x8_KV2x4x16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_float16_128x128x128_Q2x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x4x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x2x16_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x4x16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x4x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x2x16_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x4x16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x4x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x2x16_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x4x16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_9.cu b/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_9.cu new file mode 100644 index 00000000..b867cdc0 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/blackwell_fna_bwd/source_9.cu @@ -0,0 +1,2758 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_BLACKWELL_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_blackwell { + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x4x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x2x16_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x4x16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x4x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x2x16_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x4x16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x4x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x2x16_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x4x16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q1x8x16_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x4x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x2x16_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<2>, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x4x16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<16>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q2x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void blackwell_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using Config = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/fmha/source_0.cu b/natten/csrc/autogen/src/cuda/fmha/source_0.cu new file mode 100644 index 00000000..cfe0a869 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fmha/source_0.cu @@ -0,0 +1,780 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha { + + + +/////////////////////////////////////////////////////////////////// +// FMHA / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x32_sm50_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x32_sm50_float32` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x32_sm50_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x32_sm50_float32` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x32_sm50_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x32_sm50_float32` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x64_sm50_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x64_sm50_float32` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x64_sm50_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x64_sm50_float32` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x64_sm50_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x64_sm50_float32` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x128_sm50_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x128_sm50_float32` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x128_sm50_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x128_sm50_float32` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x128_sm50_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x128_sm50_float32` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x65536_sm50_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x65536_sm50_float32` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x65536_sm50_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x65536_sm50_float32` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x65536_sm50_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x65536_sm50_float32` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FMHA / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x32_sm50_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x32_sm50_float16` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x32_sm50_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x32_sm50_float16` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x32_sm50_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x32_sm50_float16` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x64_sm50_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x64_sm50_float16` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x64_sm50_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x64_sm50_float16` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x64_sm50_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x64_sm50_float16` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x128_sm50_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x128_sm50_float16` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x128_sm50_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x128_sm50_float16` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x128_sm50_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x128_sm50_float16` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x65536_sm50_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x65536_sm50_float16` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x65536_sm50_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x65536_sm50_float16` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x65536_sm50_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x65536_sm50_float16` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FMHA / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x32_sm70_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x32_sm70_float32` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x32_sm70_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x32_sm70_float32` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x32_sm70_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x32_sm70_float32` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x64_sm70_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x64_sm70_float32` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x64_sm70_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x64_sm70_float32` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x64_sm70_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x64_sm70_float32` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x128_sm70_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x128_sm70_float32` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x128_sm70_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x128_sm70_float32` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x128_sm70_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x128_sm70_float32` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x65536_sm70_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x65536_sm70_float32` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x65536_sm70_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x65536_sm70_float32` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x65536_sm70_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x65536_sm70_float32` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fmha/source_1.cu b/natten/csrc/autogen/src/cuda/fmha/source_1.cu new file mode 100644 index 00000000..cdaa56d9 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fmha/source_1.cu @@ -0,0 +1,780 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha { + + + +/////////////////////////////////////////////////////////////////// +// FMHA / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x32_sm70_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x32_sm70_float16` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x32_sm70_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x32_sm70_float16` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x32_sm70_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x32_sm70_float16` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x64_sm70_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x64_sm70_float16` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x64_sm70_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x64_sm70_float16` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x64_sm70_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x64_sm70_float16` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x128_sm70_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x128_sm70_float16` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x128_sm70_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x128_sm70_float16` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x128_sm70_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x128_sm70_float16` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x65536_sm70_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x65536_sm70_float16` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x65536_sm70_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x65536_sm70_float16` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x65536_sm70_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x65536_sm70_float16` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FMHA / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x32_sm75_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x32_sm75_float32` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x32_sm75_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x32_sm75_float32` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x32_sm75_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x32_sm75_float32` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x64_sm75_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x64_sm75_float32` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x64_sm75_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x64_sm75_float32` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x64_sm75_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x64_sm75_float32` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x128_sm75_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x128_sm75_float32` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x128_sm75_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x128_sm75_float32` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x128_sm75_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x128_sm75_float32` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x65536_sm75_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x65536_sm75_float32` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x65536_sm75_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x65536_sm75_float32` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x65536_sm75_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x65536_sm75_float32` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FMHA / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x32_sm75_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x32_sm75_float16` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x32_sm75_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x32_sm75_float16` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x32_sm75_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x32_sm75_float16` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x64_sm75_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x64_sm75_float16` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x64_sm75_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x64_sm75_float16` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x64_sm75_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x64_sm75_float16` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x128_sm75_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x128_sm75_float16` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x128_sm75_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x128_sm75_float16` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x128_sm75_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x128_sm75_float16` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x65536_sm75_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x65536_sm75_float16` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x65536_sm75_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x65536_sm75_float16` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x65536_sm75_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x65536_sm75_float16` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fmha/source_2.cu b/natten/csrc/autogen/src/cuda/fmha/source_2.cu new file mode 100644 index 00000000..77459606 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fmha/source_2.cu @@ -0,0 +1,780 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha { + + + +/////////////////////////////////////////////////////////////////// +// FMHA / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x32_sm80_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x32_sm80_float32` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x32_sm80_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x32_sm80_float32` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x32_sm80_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x32_sm80_float32` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x64_sm80_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x64_sm80_float32` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x64_sm80_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x64_sm80_float32` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x64_sm80_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x64_sm80_float32` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x128_sm80_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x128_sm80_float32` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x128_sm80_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x128_sm80_float32` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x128_sm80_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x128_sm80_float32` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x65536_sm80_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x65536_sm80_float32` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x65536_sm80_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x65536_sm80_float32` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x65536_sm80_float32(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x65536_sm80_float32` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FMHA / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x32_sm80_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x32_sm80_float16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x32_sm80_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x32_sm80_float16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x32_sm80_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x32_sm80_float16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x64_sm80_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x64_sm80_float16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x64_sm80_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x64_sm80_float16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x64_sm80_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x64_sm80_float16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x128_sm80_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x128_sm80_float16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x128_sm80_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x128_sm80_float16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x128_sm80_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x128_sm80_float16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x65536_sm80_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x65536_sm80_float16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x65536_sm80_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x65536_sm80_float16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x65536_sm80_float16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x65536_sm80_float16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FMHA / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x32_sm80_bfloat16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x32_sm80_bfloat16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x32_sm80_bfloat16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x32_sm80_bfloat16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x32_sm80_bfloat16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x32_sm80_bfloat16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x64_sm80_bfloat16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x64_sm80_bfloat16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x64_sm80_bfloat16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x64_sm80_bfloat16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x64_sm80_bfloat16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x64_sm80_bfloat16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x128_sm80_bfloat16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x128_sm80_bfloat16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x128_sm80_bfloat16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x128_sm80_bfloat16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x128_sm80_bfloat16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x128_sm80_bfloat16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_32x128x65536_sm80_bfloat16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_32x128x65536_sm80_bfloat16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x128x65536_sm80_bfloat16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x128x65536_sm80_bfloat16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionKernel::kNumThreads, + AttentionKernel::kMinBlocksPerSm) +fmha_64x64x65536_sm80_bfloat16(typename AttentionKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_64x64x65536_sm80_bfloat16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fmha/source_3.cu b/natten/csrc/autogen/src/cuda/fmha/source_3.cu new file mode 100644 index 00000000..61c4bd42 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fmha/source_3.cu @@ -0,0 +1,276 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha { + + + +/////////////////////////////////////////////////////////////////// +// FMHA / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x32_sm50_float32(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x32_sm50_float32` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x64_sm50_float32(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x64_sm50_float32` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x128_sm50_float32(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x128_sm50_float32` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x65536_sm50_float32(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x65536_sm50_float32` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FMHA / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x32_sm50_float16(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x32_sm50_float16` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x64_sm50_float16(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x64_sm50_float16` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x128_sm50_float16(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x128_sm50_float16` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x65536_sm50_float16(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x65536_sm50_float16` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FMHA / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x32_sm70_float32(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x32_sm70_float32` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x64_sm70_float32(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x64_sm70_float32` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x128_sm70_float32(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x128_sm70_float32` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x65536_sm70_float32(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x65536_sm70_float32` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fmha/source_4.cu b/natten/csrc/autogen/src/cuda/fmha/source_4.cu new file mode 100644 index 00000000..9fd6bdcf --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fmha/source_4.cu @@ -0,0 +1,318 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha { + + + +/////////////////////////////////////////////////////////////////// +// FMHA / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x32_sm70_float16(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x32_sm70_float16` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x64_sm70_float16(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x64_sm70_float16` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x128_sm70_float16(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x128_sm70_float16` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_128x64x128_sm70_float16(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_128x64x128_sm70_float16` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x65536_sm70_float16(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x65536_sm70_float16` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_128x64x65536_sm70_float16(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_128x64x65536_sm70_float16` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FMHA / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x32_sm75_float32(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x32_sm75_float32` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x64_sm75_float32(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x64_sm75_float32` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x128_sm75_float32(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x128_sm75_float32` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x65536_sm75_float32(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x65536_sm75_float32` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FMHA / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x32_sm75_float16(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x32_sm75_float16` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x64_sm75_float16(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x64_sm75_float16` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x128_sm75_float16(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x128_sm75_float16` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x65536_sm75_float16(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x65536_sm75_float16` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fmha/source_5.cu b/natten/csrc/autogen/src/cuda/fmha/source_5.cu new file mode 100644 index 00000000..6e8f4c77 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fmha/source_5.cu @@ -0,0 +1,402 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha { + + + +/////////////////////////////////////////////////////////////////// +// FMHA / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x32_sm80_float32(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x32_sm80_float32` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x64_sm80_float32(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x64_sm80_float32` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x128_sm80_float32(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x128_sm80_float32` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_128x64x128_sm80_float32(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_128x64x128_sm80_float32` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x65536_sm80_float32(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x65536_sm80_float32` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_128x64x65536_sm80_float32(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_128x64x65536_sm80_float32` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FMHA / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x32_sm80_float16(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x32_sm80_float16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x64_sm80_float16(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x64_sm80_float16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x128_sm80_float16(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x128_sm80_float16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_128x128x128_sm80_float16(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_128x128x128_sm80_float16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x65536_sm80_float16(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x65536_sm80_float16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_128x64x65536_sm80_float16(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_128x64x65536_sm80_float16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FMHA / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x32_sm80_bfloat16(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x32_sm80_bfloat16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x64_sm80_bfloat16(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x64_sm80_bfloat16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x128_sm80_bfloat16(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x128_sm80_bfloat16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_128x128x128_sm80_bfloat16(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_128x128x128_sm80_bfloat16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_64x64x65536_sm80_bfloat16(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_64x64x65536_sm80_bfloat16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + AttentionBackwardKernel::kNumThreads, + AttentionBackwardKernel::kMinBlocksPerSm) +fmha_backward_128x64x65536_sm80_bfloat16(typename AttentionBackwardKernel::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + AttentionBackwardKernel::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `fmha_backward_128x64x65536_sm80_bfloat16` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_0.cu b/natten/csrc/autogen/src/cuda/fna/source_0.cu new file mode 100644 index 00000000..ee1f06a4 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_0.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm50_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x32_sm50_float32_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm50_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x32_sm50_float32_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm50_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x32_sm50_float32_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm50_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x64_sm50_float32_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm50_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x64_sm50_float32_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm50_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x64_sm50_float32_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm50_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x128_sm50_float32_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm50_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x128_sm50_float32_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm50_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x128_sm50_float32_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm50_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x65536_sm50_float32_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm50_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x65536_sm50_float32_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm50_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x65536_sm50_float32_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm50_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x32_sm50_float32_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm50_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x32_sm50_float32_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm50_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x32_sm50_float32_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm50_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x64_sm50_float32_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm50_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x64_sm50_float32_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm50_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x64_sm50_float32_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm50_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x128_sm50_float32_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm50_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x128_sm50_float32_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm50_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x128_sm50_float32_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm50_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x65536_sm50_float32_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm50_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x65536_sm50_float32_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm50_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x65536_sm50_float32_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm50_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x32_sm50_float16_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm50_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x32_sm50_float16_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm50_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x32_sm50_float16_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm50_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x64_sm50_float16_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm50_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x64_sm50_float16_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm50_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x64_sm50_float16_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm50_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x128_sm50_float16_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm50_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x128_sm50_float16_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm50_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x128_sm50_float16_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm50_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x65536_sm50_float16_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm50_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x65536_sm50_float16_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm50_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x65536_sm50_float16_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm50_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x32_sm50_float16_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm50_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x32_sm50_float16_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm50_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x32_sm50_float16_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm50_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x64_sm50_float16_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm50_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x64_sm50_float16_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm50_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x64_sm50_float16_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm50_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x128_sm50_float16_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm50_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x128_sm50_float16_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm50_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x128_sm50_float16_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm50_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x65536_sm50_float16_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm50_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x65536_sm50_float16_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm50_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x65536_sm50_float16_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_1.cu b/natten/csrc/autogen/src/cuda/fna/source_1.cu new file mode 100644 index 00000000..845b7f0a --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_1.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm70_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x32_sm70_float32_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm70_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x32_sm70_float32_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm70_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x32_sm70_float32_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm70_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x64_sm70_float32_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm70_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x64_sm70_float32_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm70_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x64_sm70_float32_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm70_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x128_sm70_float32_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm70_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x128_sm70_float32_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm70_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x128_sm70_float32_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm70_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x65536_sm70_float32_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm70_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x65536_sm70_float32_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm70_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x65536_sm70_float32_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm70_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x32_sm70_float32_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm70_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x32_sm70_float32_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm70_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x32_sm70_float32_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm70_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x64_sm70_float32_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm70_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x64_sm70_float32_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm70_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x64_sm70_float32_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm70_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x128_sm70_float32_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm70_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x128_sm70_float32_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm70_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x128_sm70_float32_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm70_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x65536_sm70_float32_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm70_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x65536_sm70_float32_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm70_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x65536_sm70_float32_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm70_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x32_sm70_float16_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm70_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x32_sm70_float16_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm70_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x32_sm70_float16_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm70_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x64_sm70_float16_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm70_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x64_sm70_float16_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm70_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x64_sm70_float16_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm70_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x128_sm70_float16_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm70_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x128_sm70_float16_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm70_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x128_sm70_float16_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm70_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x65536_sm70_float16_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm70_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x65536_sm70_float16_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm70_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x65536_sm70_float16_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm70_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x32_sm70_float16_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm70_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x32_sm70_float16_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm70_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x32_sm70_float16_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm70_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x64_sm70_float16_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm70_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x64_sm70_float16_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm70_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x64_sm70_float16_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm70_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x128_sm70_float16_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm70_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x128_sm70_float16_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm70_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x128_sm70_float16_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm70_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x65536_sm70_float16_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm70_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x65536_sm70_float16_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm70_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x65536_sm70_float16_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_10.cu b/natten/csrc/autogen/src/cuda/fna/source_10.cu new file mode 100644 index 00000000..e75f68b3 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_10.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm75_float16_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm75_float16_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm75_float16_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm75_float16_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm75_float16_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm75_float16_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm75_float16_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm75_float16_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm75_float16_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm75_float16_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm75_float16_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm75_float16_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm75_float16_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm75_float16_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm75_float16_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm75_float16_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm75_float16_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm75_float16_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm75_float16_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm75_float16_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm75_float16_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm75_float16_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm75_float16_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm75_float16_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm80_float32_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm80_float32_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm80_float32_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm80_float32_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm80_float32_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm80_float32_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm80_float32_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm80_float32_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm80_float32_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm80_float32_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm80_float32_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm80_float32_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm80_float32_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm80_float32_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm80_float32_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm80_float32_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm80_float32_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm80_float32_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm80_float32_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm80_float32_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm80_float32_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm80_float32_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm80_float32_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm80_float32_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_11.cu b/natten/csrc/autogen/src/cuda/fna/source_11.cu new file mode 100644 index 00000000..e53694a2 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_11.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm80_float32_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm80_float32_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm80_float32_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm80_float32_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm80_float32_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm80_float32_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm80_float32_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm80_float32_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm80_float32_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm80_float32_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm80_float32_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm80_float32_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm80_float32_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm80_float32_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm80_float32_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm80_float32_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm80_float32_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm80_float32_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm80_float32_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm80_float32_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm80_float32_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm80_float32_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm80_float32_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm80_float32_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm80_float16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm80_float16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm80_float16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm80_float16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm80_float16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm80_float16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm80_float16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm80_float16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm80_float16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm80_float16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm80_float16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm80_float16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm80_float16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm80_float16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm80_float16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm80_float16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm80_float16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm80_float16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm80_float16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm80_float16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm80_float16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm80_float16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm80_float16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm80_float16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_12.cu b/natten/csrc/autogen/src/cuda/fna/source_12.cu new file mode 100644 index 00000000..23e48e6e --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_12.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm80_float16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm80_float16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm80_float16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm80_float16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm80_float16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm80_float16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm80_float16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm80_float16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm80_float16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm80_float16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm80_float16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm80_float16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm80_float16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm80_float16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm80_float16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm80_float16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm80_float16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm80_float16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm80_float16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm80_float16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm80_float16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm80_float16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm80_float16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm80_float16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm80_bfloat16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm80_bfloat16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm80_bfloat16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm80_bfloat16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm80_bfloat16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm80_bfloat16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm80_bfloat16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm80_bfloat16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm80_bfloat16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm80_bfloat16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm80_bfloat16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm80_bfloat16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm80_bfloat16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm80_bfloat16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm80_bfloat16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm80_bfloat16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm80_bfloat16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm80_bfloat16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm80_bfloat16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm80_bfloat16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm80_bfloat16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm80_bfloat16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm80_bfloat16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm80_bfloat16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_13.cu b/natten/csrc/autogen/src/cuda/fna/source_13.cu new file mode 100644 index 00000000..22e9bec6 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_13.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm80_bfloat16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm80_bfloat16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm80_bfloat16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm80_bfloat16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm80_bfloat16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm80_bfloat16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm80_bfloat16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm80_bfloat16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm80_bfloat16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm80_bfloat16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm80_bfloat16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm80_bfloat16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm80_bfloat16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm80_bfloat16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm80_bfloat16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm80_bfloat16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm80_bfloat16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm80_bfloat16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm80_bfloat16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm80_bfloat16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm80_bfloat16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm80_bfloat16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm80_bfloat16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm80_bfloat16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm50_float32_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm50_float32_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm50_float32_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm50_float32_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm50_float32_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm50_float32_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm50_float32_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm50_float32_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm50_float32_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm50_float32_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm50_float32_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm50_float32_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm50_float32_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm50_float32_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm50_float32_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm50_float32_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm50_float32_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm50_float32_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm50_float32_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm50_float32_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm50_float32_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm50_float32_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm50_float32_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm50_float32_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_14.cu b/natten/csrc/autogen/src/cuda/fna/source_14.cu new file mode 100644 index 00000000..ff1558b2 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_14.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm50_float32_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm50_float32_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm50_float32_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm50_float32_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm50_float32_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm50_float32_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm50_float32_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm50_float32_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm50_float32_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm50_float32_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm50_float32_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm50_float32_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm50_float32_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm50_float32_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm50_float32_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm50_float32_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm50_float32_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm50_float32_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm50_float32_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm50_float32_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm50_float32_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm50_float32_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm50_float32_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm50_float32_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm50_float32_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm50_float32_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm50_float32_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm50_float32_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm50_float32_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm50_float32_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm50_float32_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm50_float32_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm50_float32_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm50_float32_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm50_float32_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm50_float32_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm50_float32_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm50_float32_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm50_float32_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm50_float32_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm50_float32_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm50_float32_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm50_float32_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm50_float32_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm50_float32_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm50_float32_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm50_float32_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm50_float32_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_15.cu b/natten/csrc/autogen/src/cuda/fna/source_15.cu new file mode 100644 index 00000000..05b247b7 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_15.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm50_float32_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm50_float32_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm50_float32_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm50_float32_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm50_float32_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm50_float32_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm50_float32_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm50_float32_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm50_float32_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm50_float32_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm50_float32_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm50_float32_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm50_float32_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm50_float32_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm50_float32_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm50_float32_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm50_float32_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm50_float32_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm50_float32_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm50_float32_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm50_float32_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm50_float32_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm50_float32_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm50_float32_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm50_float16_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm50_float16_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm50_float16_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm50_float16_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm50_float16_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm50_float16_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm50_float16_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm50_float16_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm50_float16_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm50_float16_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm50_float16_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm50_float16_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm50_float16_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm50_float16_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm50_float16_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm50_float16_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm50_float16_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm50_float16_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm50_float16_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm50_float16_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm50_float16_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm50_float16_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm50_float16_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm50_float16_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_16.cu b/natten/csrc/autogen/src/cuda/fna/source_16.cu new file mode 100644 index 00000000..fcf0232d --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_16.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm50_float16_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm50_float16_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm50_float16_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm50_float16_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm50_float16_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm50_float16_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm50_float16_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm50_float16_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm50_float16_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm50_float16_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm50_float16_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm50_float16_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm50_float16_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm50_float16_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm50_float16_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm50_float16_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm50_float16_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm50_float16_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm50_float16_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm50_float16_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm50_float16_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm50_float16_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm50_float16_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm50_float16_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm50_float16_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm50_float16_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm50_float16_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm50_float16_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm50_float16_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm50_float16_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm50_float16_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm50_float16_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm50_float16_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm50_float16_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm50_float16_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm50_float16_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm50_float16_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm50_float16_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm50_float16_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm50_float16_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm50_float16_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm50_float16_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm50_float16_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm50_float16_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm50_float16_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm50_float16_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm50_float16_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm50_float16_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_17.cu b/natten/csrc/autogen/src/cuda/fna/source_17.cu new file mode 100644 index 00000000..fbb31877 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_17.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm50_float16_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm50_float16_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm50_float16_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm50_float16_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm50_float16_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm50_float16_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm50_float16_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm50_float16_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm50_float16_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm50_float16_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm50_float16_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm50_float16_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm50_float16_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm50_float16_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm50_float16_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm50_float16_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm50_float16_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm50_float16_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm50_float16_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm50_float16_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm50_float16_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm50_float16_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm50_float16_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm50_float16_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm70_float32_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm70_float32_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm70_float32_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm70_float32_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm70_float32_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm70_float32_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm70_float32_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm70_float32_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm70_float32_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm70_float32_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm70_float32_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm70_float32_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm70_float32_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm70_float32_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm70_float32_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm70_float32_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm70_float32_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm70_float32_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm70_float32_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm70_float32_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm70_float32_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm70_float32_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm70_float32_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm70_float32_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_18.cu b/natten/csrc/autogen/src/cuda/fna/source_18.cu new file mode 100644 index 00000000..e8f378f9 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_18.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm70_float32_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm70_float32_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm70_float32_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm70_float32_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm70_float32_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm70_float32_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm70_float32_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm70_float32_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm70_float32_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm70_float32_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm70_float32_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm70_float32_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm70_float32_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm70_float32_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm70_float32_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm70_float32_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm70_float32_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm70_float32_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm70_float32_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm70_float32_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm70_float32_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm70_float32_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm70_float32_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm70_float32_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm70_float32_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm70_float32_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm70_float32_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm70_float32_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm70_float32_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm70_float32_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm70_float32_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm70_float32_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm70_float32_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm70_float32_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm70_float32_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm70_float32_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm70_float32_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm70_float32_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm70_float32_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm70_float32_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm70_float32_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm70_float32_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm70_float32_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm70_float32_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm70_float32_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm70_float32_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm70_float32_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm70_float32_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_19.cu b/natten/csrc/autogen/src/cuda/fna/source_19.cu new file mode 100644 index 00000000..0e084d99 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_19.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm70_float32_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm70_float32_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm70_float32_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm70_float32_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm70_float32_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm70_float32_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm70_float32_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm70_float32_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm70_float32_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm70_float32_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm70_float32_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm70_float32_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm70_float32_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm70_float32_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm70_float32_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm70_float32_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm70_float32_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm70_float32_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm70_float32_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm70_float32_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm70_float32_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm70_float32_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm70_float32_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm70_float32_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm70_float16_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm70_float16_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm70_float16_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm70_float16_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm70_float16_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm70_float16_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm70_float16_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm70_float16_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm70_float16_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm70_float16_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm70_float16_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm70_float16_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm70_float16_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm70_float16_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm70_float16_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm70_float16_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm70_float16_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm70_float16_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm70_float16_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm70_float16_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm70_float16_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm70_float16_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm70_float16_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm70_float16_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_2.cu b/natten/csrc/autogen/src/cuda/fna/source_2.cu new file mode 100644 index 00000000..acb1e745 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_2.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm75_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x32_sm75_float32_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm75_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x32_sm75_float32_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm75_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x32_sm75_float32_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm75_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x64_sm75_float32_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm75_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x64_sm75_float32_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm75_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x64_sm75_float32_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm75_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x128_sm75_float32_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm75_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x128_sm75_float32_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm75_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x128_sm75_float32_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm75_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x65536_sm75_float32_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm75_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x65536_sm75_float32_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm75_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x65536_sm75_float32_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm75_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x32_sm75_float32_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm75_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x32_sm75_float32_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm75_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x32_sm75_float32_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm75_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x64_sm75_float32_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm75_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x64_sm75_float32_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm75_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x64_sm75_float32_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm75_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x128_sm75_float32_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm75_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x128_sm75_float32_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm75_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x128_sm75_float32_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm75_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x65536_sm75_float32_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm75_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x65536_sm75_float32_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm75_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x65536_sm75_float32_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm75_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x32_sm75_float16_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm75_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x32_sm75_float16_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm75_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x32_sm75_float16_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm75_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x64_sm75_float16_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm75_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x64_sm75_float16_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm75_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x64_sm75_float16_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm75_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x128_sm75_float16_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm75_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x128_sm75_float16_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm75_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x128_sm75_float16_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm75_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x65536_sm75_float16_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm75_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x65536_sm75_float16_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm75_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x65536_sm75_float16_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm75_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x32_sm75_float16_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm75_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x32_sm75_float16_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm75_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x32_sm75_float16_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm75_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x64_sm75_float16_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm75_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x64_sm75_float16_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm75_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x64_sm75_float16_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm75_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x128_sm75_float16_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm75_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x128_sm75_float16_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm75_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x128_sm75_float16_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm75_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x65536_sm75_float16_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm75_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x65536_sm75_float16_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm75_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x65536_sm75_float16_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_20.cu b/natten/csrc/autogen/src/cuda/fna/source_20.cu new file mode 100644 index 00000000..cd9c12eb --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_20.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm70_float16_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm70_float16_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm70_float16_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm70_float16_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm70_float16_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm70_float16_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm70_float16_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm70_float16_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm70_float16_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm70_float16_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm70_float16_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm70_float16_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm70_float16_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm70_float16_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm70_float16_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm70_float16_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm70_float16_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm70_float16_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm70_float16_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm70_float16_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm70_float16_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm70_float16_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm70_float16_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm70_float16_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm70_float16_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm70_float16_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm70_float16_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm70_float16_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm70_float16_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm70_float16_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm70_float16_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm70_float16_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm70_float16_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm70_float16_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm70_float16_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm70_float16_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm70_float16_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm70_float16_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm70_float16_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm70_float16_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm70_float16_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm70_float16_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm70_float16_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm70_float16_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm70_float16_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm70_float16_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm70_float16_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm70_float16_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_21.cu b/natten/csrc/autogen/src/cuda/fna/source_21.cu new file mode 100644 index 00000000..10ababb9 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_21.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm70_float16_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm70_float16_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm70_float16_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm70_float16_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm70_float16_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm70_float16_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm70_float16_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm70_float16_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm70_float16_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm70_float16_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm70_float16_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm70_float16_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm70_float16_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm70_float16_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm70_float16_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm70_float16_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm70_float16_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm70_float16_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm70_float16_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm70_float16_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm70_float16_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm70_float16_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm70_float16_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm70_float16_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm75_float32_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm75_float32_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm75_float32_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm75_float32_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm75_float32_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm75_float32_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm75_float32_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm75_float32_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm75_float32_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm75_float32_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm75_float32_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm75_float32_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm75_float32_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm75_float32_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm75_float32_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm75_float32_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm75_float32_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm75_float32_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm75_float32_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm75_float32_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm75_float32_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm75_float32_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm75_float32_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm75_float32_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_22.cu b/natten/csrc/autogen/src/cuda/fna/source_22.cu new file mode 100644 index 00000000..ab611bb6 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_22.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm75_float32_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm75_float32_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm75_float32_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm75_float32_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm75_float32_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm75_float32_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm75_float32_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm75_float32_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm75_float32_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm75_float32_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm75_float32_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm75_float32_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm75_float32_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm75_float32_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm75_float32_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm75_float32_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm75_float32_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm75_float32_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm75_float32_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm75_float32_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm75_float32_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm75_float32_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm75_float32_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm75_float32_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm75_float32_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm75_float32_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm75_float32_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm75_float32_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm75_float32_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm75_float32_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm75_float32_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm75_float32_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm75_float32_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm75_float32_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm75_float32_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm75_float32_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm75_float32_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm75_float32_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm75_float32_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm75_float32_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm75_float32_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm75_float32_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm75_float32_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm75_float32_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm75_float32_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm75_float32_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm75_float32_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm75_float32_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_23.cu b/natten/csrc/autogen/src/cuda/fna/source_23.cu new file mode 100644 index 00000000..a5214da4 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_23.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm75_float32_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm75_float32_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm75_float32_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm75_float32_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm75_float32_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm75_float32_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm75_float32_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm75_float32_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm75_float32_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm75_float32_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm75_float32_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm75_float32_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm75_float32_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm75_float32_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm75_float32_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm75_float32_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm75_float32_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm75_float32_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm75_float32_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm75_float32_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm75_float32_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm75_float32_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm75_float32_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm75_float32_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm75_float16_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm75_float16_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm75_float16_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm75_float16_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm75_float16_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm75_float16_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm75_float16_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm75_float16_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm75_float16_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm75_float16_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm75_float16_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm75_float16_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm75_float16_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm75_float16_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm75_float16_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm75_float16_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm75_float16_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm75_float16_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm75_float16_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm75_float16_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm75_float16_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm75_float16_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm75_float16_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm75_float16_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_24.cu b/natten/csrc/autogen/src/cuda/fna/source_24.cu new file mode 100644 index 00000000..f73b17d8 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_24.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm75_float16_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm75_float16_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm75_float16_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm75_float16_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm75_float16_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm75_float16_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm75_float16_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm75_float16_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm75_float16_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm75_float16_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm75_float16_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm75_float16_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm75_float16_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm75_float16_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm75_float16_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm75_float16_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm75_float16_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm75_float16_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm75_float16_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm75_float16_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm75_float16_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm75_float16_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm75_float16_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm75_float16_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm75_float16_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm75_float16_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm75_float16_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm75_float16_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm75_float16_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm75_float16_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm75_float16_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm75_float16_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm75_float16_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm75_float16_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm75_float16_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm75_float16_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm75_float16_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm75_float16_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm75_float16_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm75_float16_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm75_float16_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm75_float16_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm75_float16_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm75_float16_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm75_float16_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm75_float16_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm75_float16_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm75_float16_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_25.cu b/natten/csrc/autogen/src/cuda/fna/source_25.cu new file mode 100644 index 00000000..ba57816e --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_25.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm75_float16_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm75_float16_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm75_float16_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm75_float16_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm75_float16_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm75_float16_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm75_float16_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm75_float16_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm75_float16_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm75_float16_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm75_float16_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm75_float16_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm75_float16_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm75_float16_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm75_float16_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm75_float16_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm75_float16_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm75_float16_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm75_float16_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm75_float16_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm75_float16_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm75_float16_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm75_float16_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm75_float16_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm80_float32_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm80_float32_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm80_float32_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm80_float32_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm80_float32_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm80_float32_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm80_float32_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm80_float32_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm80_float32_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm80_float32_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm80_float32_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm80_float32_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm80_float32_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm80_float32_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm80_float32_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm80_float32_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm80_float32_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm80_float32_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm80_float32_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm80_float32_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm80_float32_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm80_float32_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm80_float32_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm80_float32_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_26.cu b/natten/csrc/autogen/src/cuda/fna/source_26.cu new file mode 100644 index 00000000..b821660f --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_26.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm80_float32_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm80_float32_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm80_float32_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm80_float32_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm80_float32_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm80_float32_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm80_float32_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm80_float32_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm80_float32_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm80_float32_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm80_float32_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm80_float32_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm80_float32_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm80_float32_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm80_float32_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm80_float32_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm80_float32_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm80_float32_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm80_float32_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm80_float32_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm80_float32_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm80_float32_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm80_float32_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm80_float32_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm80_float32_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm80_float32_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm80_float32_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm80_float32_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm80_float32_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm80_float32_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm80_float32_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm80_float32_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm80_float32_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm80_float32_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm80_float32_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm80_float32_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm80_float32_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm80_float32_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm80_float32_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm80_float32_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm80_float32_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm80_float32_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm80_float32_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm80_float32_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm80_float32_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm80_float32_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm80_float32_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm80_float32_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_27.cu b/natten/csrc/autogen/src/cuda/fna/source_27.cu new file mode 100644 index 00000000..43dad400 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_27.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm80_float32_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm80_float32_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm80_float32_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm80_float32_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm80_float32_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm80_float32_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm80_float32_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm80_float32_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm80_float32_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm80_float32_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm80_float32_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm80_float32_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm80_float32_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm80_float32_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm80_float32_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm80_float32_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm80_float32_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm80_float32_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm80_float32_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm80_float32_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm80_float32_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm80_float32_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm80_float32_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm80_float32_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm80_float16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm80_float16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm80_float16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm80_float16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm80_float16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm80_float16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm80_float16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm80_float16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm80_float16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm80_float16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm80_float16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm80_float16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm80_float16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm80_float16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm80_float16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm80_float16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm80_float16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm80_float16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm80_float16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm80_float16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm80_float16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm80_float16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm80_float16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm80_float16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_28.cu b/natten/csrc/autogen/src/cuda/fna/source_28.cu new file mode 100644 index 00000000..471a5600 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_28.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm80_float16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm80_float16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm80_float16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm80_float16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm80_float16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm80_float16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm80_float16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm80_float16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm80_float16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm80_float16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm80_float16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm80_float16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm80_float16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm80_float16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm80_float16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm80_float16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm80_float16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm80_float16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm80_float16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm80_float16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm80_float16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm80_float16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm80_float16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm80_float16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm80_float16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm80_float16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm80_float16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm80_float16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm80_float16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm80_float16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm80_float16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm80_float16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm80_float16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm80_float16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm80_float16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm80_float16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm80_float16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm80_float16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm80_float16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm80_float16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm80_float16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm80_float16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm80_float16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm80_float16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm80_float16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm80_float16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm80_float16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm80_float16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_29.cu b/natten/csrc/autogen/src/cuda/fna/source_29.cu new file mode 100644 index 00000000..12e66518 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_29.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm80_float16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm80_float16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm80_float16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm80_float16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm80_float16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm80_float16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm80_float16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm80_float16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm80_float16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm80_float16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm80_float16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm80_float16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm80_float16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm80_float16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm80_float16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm80_float16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm80_float16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm80_float16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm80_float16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm80_float16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm80_float16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm80_float16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm80_float16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm80_float16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm80_bfloat16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm80_bfloat16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm80_bfloat16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm80_bfloat16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm80_bfloat16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm80_bfloat16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm80_bfloat16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm80_bfloat16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm80_bfloat16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm80_bfloat16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm80_bfloat16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm80_bfloat16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm80_bfloat16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm80_bfloat16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm80_bfloat16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm80_bfloat16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm80_bfloat16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm80_bfloat16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm80_bfloat16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm80_bfloat16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm80_bfloat16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm80_bfloat16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm80_bfloat16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm80_bfloat16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_3.cu b/natten/csrc/autogen/src/cuda/fna/source_3.cu new file mode 100644 index 00000000..04f2beed --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_3.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm80_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x32_sm80_float32_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm80_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x32_sm80_float32_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm80_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x32_sm80_float32_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm80_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x64_sm80_float32_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm80_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x64_sm80_float32_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm80_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x64_sm80_float32_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm80_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x128_sm80_float32_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm80_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x128_sm80_float32_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm80_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x128_sm80_float32_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm80_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x65536_sm80_float32_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm80_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x65536_sm80_float32_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm80_float32_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x65536_sm80_float32_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm80_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x32_sm80_float32_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm80_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x32_sm80_float32_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm80_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x32_sm80_float32_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm80_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x64_sm80_float32_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm80_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x64_sm80_float32_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm80_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x64_sm80_float32_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm80_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x128_sm80_float32_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm80_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x128_sm80_float32_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm80_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x128_sm80_float32_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm80_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x65536_sm80_float32_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm80_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x65536_sm80_float32_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm80_float32_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x65536_sm80_float32_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm80_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x32_sm80_float16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm80_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x32_sm80_float16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm80_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x32_sm80_float16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm80_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x64_sm80_float16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm80_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x64_sm80_float16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm80_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x64_sm80_float16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm80_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x128_sm80_float16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm80_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x128_sm80_float16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm80_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x128_sm80_float16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm80_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x65536_sm80_float16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm80_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x65536_sm80_float16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm80_float16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x65536_sm80_float16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm80_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x32_sm80_float16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm80_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x32_sm80_float16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm80_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x32_sm80_float16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm80_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x64_sm80_float16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm80_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x64_sm80_float16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm80_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x64_sm80_float16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm80_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x128_sm80_float16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm80_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x128_sm80_float16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm80_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x128_sm80_float16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm80_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x65536_sm80_float16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm80_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x65536_sm80_float16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm80_float16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x65536_sm80_float16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_30.cu b/natten/csrc/autogen/src/cuda/fna/source_30.cu new file mode 100644 index 00000000..73663bf3 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_30.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm80_bfloat16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm80_bfloat16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm80_bfloat16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm80_bfloat16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm80_bfloat16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm80_bfloat16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm80_bfloat16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm80_bfloat16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm80_bfloat16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm80_bfloat16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm80_bfloat16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm80_bfloat16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm80_bfloat16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm80_bfloat16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm80_bfloat16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm80_bfloat16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm80_bfloat16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm80_bfloat16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm80_bfloat16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm80_bfloat16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm80_bfloat16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm80_bfloat16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm80_bfloat16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm80_bfloat16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm80_bfloat16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm80_bfloat16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm80_bfloat16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm80_bfloat16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm80_bfloat16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm80_bfloat16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm80_bfloat16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm80_bfloat16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm80_bfloat16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm80_bfloat16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm80_bfloat16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm80_bfloat16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm80_bfloat16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm80_bfloat16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm80_bfloat16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm80_bfloat16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm80_bfloat16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm80_bfloat16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm80_bfloat16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm80_bfloat16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm80_bfloat16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm80_bfloat16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm80_bfloat16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm80_bfloat16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_31.cu b/natten/csrc/autogen/src/cuda/fna/source_31.cu new file mode 100644 index 00000000..86e351ed --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_31.cu @@ -0,0 +1,701 @@ +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm80_bfloat16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm80_bfloat16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm80_bfloat16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm80_bfloat16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm80_bfloat16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm80_bfloat16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm80_bfloat16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm80_bfloat16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm80_bfloat16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm80_bfloat16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm80_bfloat16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm80_bfloat16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna3d_32x128x32_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x32_sm80_bfloat16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_64x64x32_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x32_sm80_bfloat16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna3d_64x128x32_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x32_sm80_bfloat16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna3d_32x128x64_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x64_sm80_bfloat16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_64x64x64_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x64_sm80_bfloat16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna3d_64x128x64_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x64_sm80_bfloat16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna3d_32x128x128_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x128_sm80_bfloat16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_64x64x128_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x128_sm80_bfloat16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna3d_64x128x128_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x128_sm80_bfloat16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna3d_32x128x65536_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_32x128x65536_sm80_bfloat16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna3d_64x128x65536_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x128x65536_sm80_bfloat16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_64x64x65536_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_64x64x65536_sm80_bfloat16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm50_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x32_sm50_float32_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm50_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x64_sm50_float32_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm50_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x128_sm50_float32_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm50_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x65536_sm50_float32_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm50_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x32_sm50_float32_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm50_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x64_sm50_float32_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm50_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x128_sm50_float32_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm50_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x65536_sm50_float32_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_32.cu b/natten/csrc/autogen/src/cuda/fna/source_32.cu new file mode 100644 index 00000000..966ce54d --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_32.cu @@ -0,0 +1,364 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm50_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x32_sm50_float16_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm50_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x64_sm50_float16_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm50_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x128_sm50_float16_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm50_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x65536_sm50_float16_cm_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm50_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x32_sm50_float16_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm50_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x64_sm50_float16_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm50_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x128_sm50_float16_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm50_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x65536_sm50_float16_cm_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm70_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x32_sm70_float32_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm70_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x64_sm70_float32_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm70_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x128_sm70_float32_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm70_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x65536_sm70_float32_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm70_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x32_sm70_float32_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm70_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x64_sm70_float32_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm70_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x128_sm70_float32_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm70_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x65536_sm70_float32_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_33.cu b/natten/csrc/autogen/src/cuda/fna/source_33.cu new file mode 100644 index 00000000..a1e712c9 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_33.cu @@ -0,0 +1,448 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm70_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x32_sm70_float16_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm70_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x64_sm70_float16_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm70_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x128_sm70_float16_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kMinBlocksPerSm) +fna1d_backward_128x64x128_sm70_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_128x64x128_sm70_float16_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm70_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x65536_sm70_float16_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_128x64x65536_sm70_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_128x64x65536_sm70_float16_cm_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm70_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x32_sm70_float16_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm70_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x64_sm70_float16_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm70_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x128_sm70_float16_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kMinBlocksPerSm) +fna1d_backward_128x64x128_sm70_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_128x64x128_sm70_float16_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm70_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x65536_sm70_float16_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_128x64x65536_sm70_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_128x64x65536_sm70_float16_cm_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm75_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x32_sm75_float32_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm75_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x64_sm75_float32_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm75_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x128_sm75_float32_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm75_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x65536_sm75_float32_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm75_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x32_sm75_float32_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm75_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x64_sm75_float32_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm75_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x128_sm75_float32_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm75_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x65536_sm75_float32_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_34.cu b/natten/csrc/autogen/src/cuda/fna/source_34.cu new file mode 100644 index 00000000..3fe66eb2 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_34.cu @@ -0,0 +1,448 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm75_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x32_sm75_float16_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm75_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x64_sm75_float16_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm75_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x128_sm75_float16_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm75_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x65536_sm75_float16_cm_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm75_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x32_sm75_float16_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm75_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x64_sm75_float16_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm75_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x128_sm75_float16_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm75_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x65536_sm75_float16_cm_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm80_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x32_sm80_float32_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm80_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x64_sm80_float32_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm80_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x128_sm80_float32_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kMinBlocksPerSm) +fna1d_backward_128x64x128_sm80_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_128x64x128_sm80_float32_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm80_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x65536_sm80_float32_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_128x64x65536_sm80_float32_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_128x64x65536_sm80_float32_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm80_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x32_sm80_float32_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm80_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x64_sm80_float32_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm80_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x128_sm80_float32_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kMinBlocksPerSm) +fna1d_backward_128x64x128_sm80_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_128x64x128_sm80_float32_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm80_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x65536_sm80_float32_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_128x64x65536_sm80_float32_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_128x64x65536_sm80_float32_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_35.cu b/natten/csrc/autogen/src/cuda/fna/source_35.cu new file mode 100644 index 00000000..571f08b6 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_35.cu @@ -0,0 +1,532 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm80_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x32_sm80_float16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm80_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x64_sm80_float16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm80_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x128_sm80_float16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna1d_backward_128x128x128_sm80_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_128x128x128_sm80_float16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm80_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x65536_sm80_float16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_128x64x65536_sm80_float16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_128x64x65536_sm80_float16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm80_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x32_sm80_float16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm80_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x64_sm80_float16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm80_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x128_sm80_float16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna1d_backward_128x128x128_sm80_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_128x128x128_sm80_float16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm80_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x65536_sm80_float16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_128x64x65536_sm80_float16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_128x64x65536_sm80_float16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x32_sm80_bfloat16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x64_sm80_bfloat16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x128_sm80_bfloat16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna1d_backward_128x128x128_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_128x128x128_sm80_bfloat16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x65536_sm80_bfloat16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_128x64x65536_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_128x64x65536_sm80_bfloat16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_backward_64x64x32_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x32_sm80_bfloat16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_backward_64x64x64_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x64_sm80_bfloat16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_backward_64x64x128_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x128_sm80_bfloat16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna1d_backward_128x128x128_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_128x128x128_sm80_bfloat16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_64x64x65536_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_64x64x65536_sm80_bfloat16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna1d_backward_128x64x65536_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_backward_128x64x65536_sm80_bfloat16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_36.cu b/natten/csrc/autogen/src/cuda/fna/source_36.cu new file mode 100644 index 00000000..4d63eb86 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_36.cu @@ -0,0 +1,364 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm50_float32_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm50_float32_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm50_float32_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm50_float32_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm50_float32_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm50_float32_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm50_float32_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm50_float32_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm50_float32_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm50_float32_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm50_float32_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm50_float32_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm50_float32_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm50_float32_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm50_float32_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm50_float32_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_37.cu b/natten/csrc/autogen/src/cuda/fna/source_37.cu new file mode 100644 index 00000000..1aea629e --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_37.cu @@ -0,0 +1,364 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm50_float16_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm50_float16_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm50_float16_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm50_float16_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm50_float16_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm50_float16_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm50_float16_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm50_float16_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm50_float16_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm50_float16_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm50_float16_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm50_float16_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm50_float16_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm50_float16_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm50_float16_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm50_float16_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_38.cu b/natten/csrc/autogen/src/cuda/fna/source_38.cu new file mode 100644 index 00000000..419c96e7 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_38.cu @@ -0,0 +1,364 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm70_float32_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm70_float32_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm70_float32_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm70_float32_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm70_float32_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm70_float32_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm70_float32_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm70_float32_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm70_float32_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm70_float32_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm70_float32_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm70_float32_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm70_float32_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm70_float32_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm70_float32_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm70_float32_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_39.cu b/natten/csrc/autogen/src/cuda/fna/source_39.cu new file mode 100644 index 00000000..ff83f2b9 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_39.cu @@ -0,0 +1,532 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm70_float16_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm70_float16_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm70_float16_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kMinBlocksPerSm) +fna2d_backward_128x64x128_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x64x128_sm70_float16_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm70_float16_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x64x65536_sm70_float16_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm70_float16_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm70_float16_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm70_float16_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kMinBlocksPerSm) +fna2d_backward_128x64x128_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x64x128_sm70_float16_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm70_float16_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x64x65536_sm70_float16_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm70_float16_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm70_float16_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm70_float16_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kMinBlocksPerSm) +fna2d_backward_128x64x128_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x64x128_sm70_float16_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm70_float16_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x64x65536_sm70_float16_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm70_float16_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm70_float16_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm70_float16_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kMinBlocksPerSm) +fna2d_backward_128x64x128_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x64x128_sm70_float16_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm70_float16_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x64x65536_sm70_float16_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_4.cu b/natten/csrc/autogen/src/cuda/fna/source_4.cu new file mode 100644 index 00000000..257eb455 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_4.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x32_sm80_bfloat16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x32_sm80_bfloat16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x32_sm80_bfloat16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x64_sm80_bfloat16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x64_sm80_bfloat16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x64_sm80_bfloat16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x128_sm80_bfloat16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x128_sm80_bfloat16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x128_sm80_bfloat16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x65536_sm80_bfloat16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x65536_sm80_bfloat16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm80_bfloat16_cm_0(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x65536_sm80_bfloat16_cm_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-1D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::kMinBlocksPerSm) +fna1d_32x128x32_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x32_sm80_bfloat16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna1d_64x64x32_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x32_sm80_bfloat16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::kMinBlocksPerSm) +fna1d_64x128x32_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x32_sm80_bfloat16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::kMinBlocksPerSm) +fna1d_32x128x64_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x64_sm80_bfloat16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna1d_64x64x64_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x64_sm80_bfloat16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::kMinBlocksPerSm) +fna1d_64x128x64_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x64_sm80_bfloat16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::kMinBlocksPerSm) +fna1d_32x128x128_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x128_sm80_bfloat16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna1d_64x64x128_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x128_sm80_bfloat16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::kMinBlocksPerSm) +fna1d_64x128x128_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x128_sm80_bfloat16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::kMinBlocksPerSm) +fna1d_32x128x65536_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_32x128x65536_sm80_bfloat16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::kMinBlocksPerSm) +fna1d_64x128x65536_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x128x65536_sm80_bfloat16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna1d_64x64x65536_sm80_bfloat16_cm_1(typename FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<1, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna1d_64x64x65536_sm80_bfloat16_cm_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm50_float32_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm50_float32_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm50_float32_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm50_float32_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm50_float32_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm50_float32_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm50_float32_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm50_float32_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm50_float32_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm50_float32_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm50_float32_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm50_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm50_float32_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm50_float32_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm50_float32_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm50_float32_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm50_float32_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm50_float32_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm50_float32_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm50_float32_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm50_float32_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm50_float32_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm50_float32_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm50_float32_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm50_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm50_float32_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_40.cu b/natten/csrc/autogen/src/cuda/fna/source_40.cu new file mode 100644 index 00000000..3af9e006 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_40.cu @@ -0,0 +1,364 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm75_float32_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm75_float32_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm75_float32_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm75_float32_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm75_float32_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm75_float32_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm75_float32_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm75_float32_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm75_float32_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm75_float32_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm75_float32_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm75_float32_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm75_float32_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm75_float32_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm75_float32_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm75_float32_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_41.cu b/natten/csrc/autogen/src/cuda/fna/source_41.cu new file mode 100644 index 00000000..e098d9ce --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_41.cu @@ -0,0 +1,364 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm75_float16_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm75_float16_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm75_float16_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm75_float16_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm75_float16_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm75_float16_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm75_float16_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm75_float16_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm75_float16_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm75_float16_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm75_float16_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm75_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm75_float16_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm75_float16_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm75_float16_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm75_float16_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm75_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm75_float16_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_42.cu b/natten/csrc/autogen/src/cuda/fna/source_42.cu new file mode 100644 index 00000000..411f2123 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_42.cu @@ -0,0 +1,532 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm80_float32_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm80_float32_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm80_float32_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kMinBlocksPerSm) +fna2d_backward_128x64x128_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x64x128_sm80_float32_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm80_float32_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm80_float32_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x64x65536_sm80_float32_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm80_float32_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm80_float32_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm80_float32_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kMinBlocksPerSm) +fna2d_backward_128x64x128_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x64x128_sm80_float32_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm80_float32_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm80_float32_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x64x65536_sm80_float32_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm80_float32_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm80_float32_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm80_float32_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kMinBlocksPerSm) +fna2d_backward_128x64x128_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x64x128_sm80_float32_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm80_float32_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm80_float32_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x64x65536_sm80_float32_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm80_float32_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm80_float32_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm80_float32_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kMinBlocksPerSm) +fna2d_backward_128x64x128_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x64x128_sm80_float32_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm80_float32_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm80_float32_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x64x65536_sm80_float32_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_43.cu b/natten/csrc/autogen/src/cuda/fna/source_43.cu new file mode 100644 index 00000000..2e364fd6 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_43.cu @@ -0,0 +1,532 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm80_float16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm80_float16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm80_float16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna2d_backward_128x128x128_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x128x128_sm80_float16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm80_float16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm80_float16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x64x65536_sm80_float16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm80_float16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm80_float16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm80_float16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna2d_backward_128x128x128_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x128x128_sm80_float16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm80_float16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm80_float16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x64x65536_sm80_float16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm80_float16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm80_float16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm80_float16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna2d_backward_128x128x128_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x128x128_sm80_float16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm80_float16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm80_float16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x64x65536_sm80_float16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm80_float16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm80_float16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm80_float16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna2d_backward_128x128x128_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x128x128_sm80_float16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm80_float16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm80_float16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x64x65536_sm80_float16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_44.cu b/natten/csrc/autogen/src/cuda/fna/source_44.cu new file mode 100644 index 00000000..a78e68e5 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_44.cu @@ -0,0 +1,532 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm80_bfloat16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm80_bfloat16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm80_bfloat16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna2d_backward_128x128x128_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x128x128_sm80_bfloat16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm80_bfloat16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm80_bfloat16_cm_0_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x64x65536_sm80_bfloat16_cm_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm80_bfloat16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm80_bfloat16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm80_bfloat16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna2d_backward_128x128x128_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x128x128_sm80_bfloat16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm80_bfloat16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm80_bfloat16_cm_0_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x64x65536_sm80_bfloat16_cm_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm80_bfloat16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm80_bfloat16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm80_bfloat16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna2d_backward_128x128x128_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x128x128_sm80_bfloat16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm80_bfloat16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm80_bfloat16_cm_1_0(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x64x65536_sm80_bfloat16_cm_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_backward_64x64x32_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x32_sm80_bfloat16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_backward_64x64x64_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x64_sm80_bfloat16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_backward_64x64x128_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x128_sm80_bfloat16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna2d_backward_128x128x128_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x128x128_sm80_bfloat16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_64x64x65536_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_64x64x65536_sm80_bfloat16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna2d_backward_128x64x65536_sm80_bfloat16_cm_1_1(typename FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<2, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_backward_128x64x65536_sm80_bfloat16_cm_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_45.cu b/natten/csrc/autogen/src/cuda/fna/source_45.cu new file mode 100644 index 00000000..238758a0 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_45.cu @@ -0,0 +1,364 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm50_float32_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm50_float32_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm50_float32_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm50_float32_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm50_float32_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm50_float32_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm50_float32_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm50_float32_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm50_float32_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm50_float32_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm50_float32_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm50_float32_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm50_float32_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm50_float32_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm50_float32_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm50_float32_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_46.cu b/natten/csrc/autogen/src/cuda/fna/source_46.cu new file mode 100644 index 00000000..a5d75258 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_46.cu @@ -0,0 +1,364 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm50_float32_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm50_float32_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm50_float32_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm50_float32_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm50_float32_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm50_float32_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm50_float32_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm50_float32_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm50_float32_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm50_float32_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm50_float32_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm50_float32_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm50_float32_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm50_float32_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm50_float32_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm50_float32_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_47.cu b/natten/csrc/autogen/src/cuda/fna/source_47.cu new file mode 100644 index 00000000..ed5bb2df --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_47.cu @@ -0,0 +1,364 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm50_float16_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm50_float16_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm50_float16_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm50_float16_cm_0_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm50_float16_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm50_float16_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm50_float16_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm50_float16_cm_0_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm50_float16_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm50_float16_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm50_float16_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm50_float16_cm_0_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm50_float16_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm50_float16_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm50_float16_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm50_float16_cm_0_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_48.cu b/natten/csrc/autogen/src/cuda/fna/source_48.cu new file mode 100644 index 00000000..39438fcd --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_48.cu @@ -0,0 +1,364 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm50_float16_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm50_float16_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm50_float16_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm50_float16_cm_1_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm50_float16_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm50_float16_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm50_float16_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm50_float16_cm_1_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm50_float16_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm50_float16_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm50_float16_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm50_float16_cm_1_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm50_float16_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm50_float16_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm50_float16_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm50_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm50_float16_cm_1_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_49.cu b/natten/csrc/autogen/src/cuda/fna/source_49.cu new file mode 100644 index 00000000..368d88bc --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_49.cu @@ -0,0 +1,364 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm70_float32_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm70_float32_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm70_float32_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm70_float32_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm70_float32_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm70_float32_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm70_float32_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm70_float32_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm70_float32_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm70_float32_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm70_float32_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm70_float32_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm70_float32_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm70_float32_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm70_float32_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm70_float32_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_5.cu b/natten/csrc/autogen/src/cuda/fna/source_5.cu new file mode 100644 index 00000000..d5b32d05 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_5.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm50_float32_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm50_float32_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm50_float32_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm50_float32_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm50_float32_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm50_float32_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm50_float32_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm50_float32_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm50_float32_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm50_float32_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm50_float32_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm50_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm50_float32_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm50_float32_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm50_float32_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm50_float32_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm50_float32_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm50_float32_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm50_float32_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm50_float32_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm50_float32_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm50_float32_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm50_float32_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm50_float32_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm50_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm50_float32_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm50_float16_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm50_float16_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm50_float16_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm50_float16_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm50_float16_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm50_float16_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm50_float16_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm50_float16_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm50_float16_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm50_float16_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm50_float16_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm50_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm50_float16_cm_0_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm50_float16_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm50_float16_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm50_float16_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm50_float16_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm50_float16_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm50_float16_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm50_float16_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm50_float16_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm50_float16_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm50_float16_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm50_float16_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm50_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm50_float16_cm_0_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_50.cu b/natten/csrc/autogen/src/cuda/fna/source_50.cu new file mode 100644 index 00000000..2b3469e0 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_50.cu @@ -0,0 +1,364 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm70_float32_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm70_float32_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm70_float32_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm70_float32_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm70_float32_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm70_float32_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm70_float32_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm70_float32_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm70_float32_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm70_float32_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm70_float32_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm70_float32_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm70_float32_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm70_float32_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm70_float32_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm70_float32_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_51.cu b/natten/csrc/autogen/src/cuda/fna/source_51.cu new file mode 100644 index 00000000..582500db --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_51.cu @@ -0,0 +1,532 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm70_float16_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm70_float16_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm70_float16_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x128_sm70_float16_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm70_float16_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm70_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm70_float16_cm_0_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm70_float16_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm70_float16_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm70_float16_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x128_sm70_float16_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm70_float16_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm70_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm70_float16_cm_0_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm70_float16_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm70_float16_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm70_float16_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x128_sm70_float16_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm70_float16_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm70_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm70_float16_cm_0_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm70_float16_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm70_float16_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm70_float16_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x128_sm70_float16_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm70_float16_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm70_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm70_float16_cm_0_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_52.cu b/natten/csrc/autogen/src/cuda/fna/source_52.cu new file mode 100644 index 00000000..9b36549e --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_52.cu @@ -0,0 +1,532 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm70_float16_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm70_float16_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm70_float16_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x128_sm70_float16_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm70_float16_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm70_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm70_float16_cm_1_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm70_float16_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm70_float16_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm70_float16_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x128_sm70_float16_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm70_float16_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm70_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm70_float16_cm_1_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm70_float16_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm70_float16_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm70_float16_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x128_sm70_float16_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm70_float16_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm70_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm70_float16_cm_1_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm70_float16_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm70_float16_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm70_float16_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x128_sm70_float16_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm70_float16_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm70_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm70_float16_cm_1_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_53.cu b/natten/csrc/autogen/src/cuda/fna/source_53.cu new file mode 100644 index 00000000..41213e24 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_53.cu @@ -0,0 +1,364 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm75_float32_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm75_float32_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm75_float32_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm75_float32_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm75_float32_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm75_float32_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm75_float32_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm75_float32_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm75_float32_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm75_float32_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm75_float32_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm75_float32_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm75_float32_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm75_float32_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm75_float32_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm75_float32_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_54.cu b/natten/csrc/autogen/src/cuda/fna/source_54.cu new file mode 100644 index 00000000..77d77b93 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_54.cu @@ -0,0 +1,364 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm75_float32_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm75_float32_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm75_float32_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm75_float32_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm75_float32_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm75_float32_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm75_float32_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm75_float32_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm75_float32_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm75_float32_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm75_float32_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm75_float32_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm75_float32_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm75_float32_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm75_float32_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm75_float32_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_55.cu b/natten/csrc/autogen/src/cuda/fna/source_55.cu new file mode 100644 index 00000000..d049e2f0 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_55.cu @@ -0,0 +1,364 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm75_float16_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm75_float16_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm75_float16_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm75_float16_cm_0_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm75_float16_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm75_float16_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm75_float16_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm75_float16_cm_0_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm75_float16_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm75_float16_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm75_float16_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm75_float16_cm_0_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm75_float16_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm75_float16_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm75_float16_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm75_float16_cm_0_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_56.cu b/natten/csrc/autogen/src/cuda/fna/source_56.cu new file mode 100644 index 00000000..d880c7f4 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_56.cu @@ -0,0 +1,364 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm75_float16_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm75_float16_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm75_float16_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm75_float16_cm_1_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm75_float16_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm75_float16_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm75_float16_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm75_float16_cm_1_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm75_float16_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm75_float16_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm75_float16_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm75_float16_cm_1_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm75_float16_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm75_float16_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm75_float16_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm75_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm75_float16_cm_1_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_57.cu b/natten/csrc/autogen/src/cuda/fna/source_57.cu new file mode 100644 index 00000000..7919044d --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_57.cu @@ -0,0 +1,532 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm80_float32_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm80_float32_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm80_float32_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x128_sm80_float32_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm80_float32_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float32_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm80_float32_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm80_float32_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm80_float32_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm80_float32_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x128_sm80_float32_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm80_float32_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float32_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm80_float32_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm80_float32_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm80_float32_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm80_float32_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x128_sm80_float32_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm80_float32_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float32_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm80_float32_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm80_float32_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm80_float32_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm80_float32_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x128_sm80_float32_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm80_float32_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float32_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm80_float32_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_58.cu b/natten/csrc/autogen/src/cuda/fna/source_58.cu new file mode 100644 index 00000000..640c2be0 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_58.cu @@ -0,0 +1,532 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm80_float32_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm80_float32_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm80_float32_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x128_sm80_float32_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm80_float32_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float32_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm80_float32_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm80_float32_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm80_float32_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm80_float32_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x128_sm80_float32_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm80_float32_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float32_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm80_float32_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm80_float32_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm80_float32_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm80_float32_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x128_sm80_float32_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm80_float32_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float32_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm80_float32_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float32 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm80_float32_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm80_float32_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm80_float32_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::kMinBlocksPerSm) +fna3d_backward_128x64x128_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x128_sm80_float32_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm80_float32_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float32_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, float, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm80_float32_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_59.cu b/natten/csrc/autogen/src/cuda/fna/source_59.cu new file mode 100644 index 00000000..003a6336 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_59.cu @@ -0,0 +1,532 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm80_float16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm80_float16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm80_float16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x128x128_sm80_float16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm80_float16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm80_float16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm80_float16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm80_float16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm80_float16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x128x128_sm80_float16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm80_float16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm80_float16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm80_float16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm80_float16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm80_float16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x128x128_sm80_float16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm80_float16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm80_float16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm80_float16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm80_float16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm80_float16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x128x128_sm80_float16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm80_float16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm80_float16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_6.cu b/natten/csrc/autogen/src/cuda/fna/source_6.cu new file mode 100644 index 00000000..a48d0bce --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_6.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm50_float16_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm50_float16_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm50_float16_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm50_float16_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm50_float16_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm50_float16_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm50_float16_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm50_float16_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm50_float16_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm50_float16_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm50_float16_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm50_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm50_float16_cm_1_0` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM50 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm50_float16_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm50_float16_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm50_float16_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm50_float16_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm50_float16_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm50_float16_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm50_float16_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm50_float16_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm50_float16_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm50_float16_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm50_float16_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm50_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 500 +#if __CUDA_ARCH__ < 700 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm50, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm50_float16_cm_1_1` was built for SM50, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm70_float32_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm70_float32_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm70_float32_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm70_float32_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm70_float32_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm70_float32_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm70_float32_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm70_float32_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm70_float32_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm70_float32_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm70_float32_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm70_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm70_float32_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm70_float32_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm70_float32_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm70_float32_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm70_float32_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm70_float32_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm70_float32_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm70_float32_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm70_float32_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm70_float32_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm70_float32_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm70_float32_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm70_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm70_float32_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_60.cu b/natten/csrc/autogen/src/cuda/fna/source_60.cu new file mode 100644 index 00000000..59181a13 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_60.cu @@ -0,0 +1,402 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm80_float16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm80_float16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm80_float16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x128x128_sm80_float16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm80_float16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm80_float16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm80_float16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm80_float16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm80_float16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x128x128_sm80_float16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm80_float16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm80_float16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm80_float16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm80_float16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm80_float16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x128x128_sm80_float16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm80_float16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm80_float16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_61.cu b/natten/csrc/autogen/src/cuda/fna/source_61.cu new file mode 100644 index 00000000..6897f6e2 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_61.cu @@ -0,0 +1,402 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / float16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm80_float16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm80_float16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm80_float16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x128x128_sm80_float16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm80_float16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_float16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::half_t, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm80_float16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm80_bfloat16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm80_bfloat16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm80_bfloat16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x128x128_sm80_bfloat16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm80_bfloat16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_bfloat16_cm_0_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm80_bfloat16_cm_0_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm80_bfloat16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm80_bfloat16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm80_bfloat16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x128x128_sm80_bfloat16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm80_bfloat16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_bfloat16_cm_0_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm80_bfloat16_cm_0_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_62.cu b/natten/csrc/autogen/src/cuda/fna/source_62.cu new file mode 100644 index 00000000..7f273026 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_62.cu @@ -0,0 +1,402 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm80_bfloat16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm80_bfloat16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm80_bfloat16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x128x128_sm80_bfloat16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm80_bfloat16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_bfloat16_cm_0_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm80_bfloat16_cm_0_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm80_bfloat16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm80_bfloat16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm80_bfloat16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x128x128_sm80_bfloat16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm80_bfloat16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_bfloat16_cm_0_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm80_bfloat16_cm_0_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm80_bfloat16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm80_bfloat16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm80_bfloat16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x128x128_sm80_bfloat16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm80_bfloat16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_bfloat16_cm_1_0_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm80_bfloat16_cm_1_0_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_63.cu b/natten/csrc/autogen/src/cuda/fna/source_63.cu new file mode 100644 index 00000000..db17f8e2 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_63.cu @@ -0,0 +1,402 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm80_bfloat16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm80_bfloat16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm80_bfloat16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x128x128_sm80_bfloat16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm80_bfloat16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_bfloat16_cm_1_0_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm80_bfloat16_cm_1_0_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm80_bfloat16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm80_bfloat16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm80_bfloat16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x128x128_sm80_bfloat16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm80_bfloat16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_bfloat16_cm_1_1_0(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm80_bfloat16_cm_1_1_0` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-3D / bfloat16 / SM80 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::kMinBlocksPerSm) +fna3d_backward_64x64x32_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x32_sm80_bfloat16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::kMinBlocksPerSm) +fna3d_backward_64x64x64_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x64_sm80_bfloat16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::kMinBlocksPerSm) +fna3d_backward_64x64x128_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x128_sm80_bfloat16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::kMinBlocksPerSm) +fna3d_backward_128x128x128_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x128x128_sm80_bfloat16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_64x64x65536_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_64x64x65536_sm80_bfloat16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::kMinBlocksPerSm) +fna3d_backward_128x64x65536_sm80_bfloat16_cm_1_1_1(typename FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#if __CUDA_ARCH__ < 1300 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionBackwardKernel<3, CausalMask, cutlass::bfloat16_t, cutlass::arch::Sm80, true, 128, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna3d_backward_128x64x65536_sm80_bfloat16_cm_1_1_1` was built for SM80, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_7.cu b/natten/csrc/autogen/src/cuda/fna/source_7.cu new file mode 100644 index 00000000..19277ac8 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_7.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm70_float32_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm70_float32_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm70_float32_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm70_float32_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm70_float32_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm70_float32_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm70_float32_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm70_float32_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm70_float32_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm70_float32_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm70_float32_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm70_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm70_float32_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm70_float32_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm70_float32_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm70_float32_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm70_float32_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm70_float32_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm70_float32_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm70_float32_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm70_float32_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm70_float32_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm70_float32_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm70_float32_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm70_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm70_float32_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm70_float16_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm70_float16_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm70_float16_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm70_float16_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm70_float16_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm70_float16_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm70_float16_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm70_float16_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm70_float16_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm70_float16_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm70_float16_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm70_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm70_float16_cm_0_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm70_float16_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm70_float16_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm70_float16_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm70_float16_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm70_float16_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm70_float16_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm70_float16_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm70_float16_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm70_float16_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm70_float16_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm70_float16_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm70_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm70_float16_cm_0_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_8.cu b/natten/csrc/autogen/src/cuda/fna/source_8.cu new file mode 100644 index 00000000..f559f50b --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_8.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm70_float16_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm70_float16_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm70_float16_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm70_float16_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm70_float16_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm70_float16_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm70_float16_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm70_float16_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm70_float16_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm70_float16_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm70_float16_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm70_float16_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm70_float16_cm_1_0` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM70 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm70_float16_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm70_float16_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm70_float16_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm70_float16_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm70_float16_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm70_float16_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm70_float16_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm70_float16_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm70_float16_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm70_float16_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm70_float16_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm70_float16_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 700 +#if __CUDA_ARCH__ < 750 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm70, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm70_float16_cm_1_1` was built for SM70, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm75_float32_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm75_float32_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm75_float32_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm75_float32_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm75_float32_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm75_float32_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm75_float32_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm75_float32_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm75_float32_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm75_float32_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm75_float32_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm75_float32_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm75_float32_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm75_float32_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm75_float32_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm75_float32_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm75_float32_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm75_float32_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm75_float32_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm75_float32_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm75_float32_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm75_float32_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm75_float32_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm75_float32_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm75_float32_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm75_float32_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/fna/source_9.cu b/natten/csrc/autogen/src/cuda/fna/source_9.cu new file mode 100644 index 00000000..bbdf937b --- /dev/null +++ b/natten/csrc/autogen/src/cuda/fna/source_9.cu @@ -0,0 +1,1036 @@ +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna { + + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm75_float32_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm75_float32_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm75_float32_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm75_float32_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm75_float32_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm75_float32_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm75_float32_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm75_float32_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm75_float32_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm75_float32_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm75_float32_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm75_float32_cm_1_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm75_float32_cm_1_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float32 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm75_float32_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm75_float32_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm75_float32_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm75_float32_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm75_float32_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm75_float32_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm75_float32_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm75_float32_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm75_float32_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm75_float32_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm75_float32_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm75_float32_cm_1_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, float, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm75_float32_cm_1_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm75_float16_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm75_float16_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm75_float16_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm75_float16_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm75_float16_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm75_float16_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm75_float16_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm75_float16_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm75_float16_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm75_float16_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm75_float16_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm75_float16_cm_0_0(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm75_float16_cm_0_0` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +/////////////////////////////////////////////////////////////////// +// FNA-2D / float16 / SM75 +/////////////////////////////////////////////////////////////////// + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::kMinBlocksPerSm) +fna2d_32x128x32_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x32_sm75_float16_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::kMinBlocksPerSm) +fna2d_64x64x32_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x32_sm75_float16_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::kMinBlocksPerSm) +fna2d_64x128x32_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 32>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x32_sm75_float16_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::kMinBlocksPerSm) +fna2d_32x128x64_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x64_sm75_float16_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::kMinBlocksPerSm) +fna2d_64x64x64_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x64_sm75_float16_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::kMinBlocksPerSm) +fna2d_64x128x64_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 64>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x64_sm75_float16_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::kMinBlocksPerSm) +fna2d_32x128x128_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x128_sm75_float16_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::kMinBlocksPerSm) +fna2d_64x64x128_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x128_sm75_float16_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::kMinBlocksPerSm) +fna2d_64x128x128_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 128>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x128_sm75_float16_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::kMinBlocksPerSm) +fna2d_32x128x65536_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 32, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_32x128x65536_sm75_float16_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::kMinBlocksPerSm) +fna2d_64x128x65536_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 128, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x128x65536_sm75_float16_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} + + +__global__ void __launch_bounds__( + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kNumThreads, + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::kMinBlocksPerSm) +fna2d_64x64x65536_sm75_float16_cm_0_1(typename FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::Params p) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 750 +#if __CUDA_ARCH__ < 800 + if (!p.advance_to_block()) { + return; + } + FusedNeighborhoodAttentionKernel<2, CausalMask, cutlass::half_t, cutlass::arch::Sm75, true, 64, 64, 65536>::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `fna2d_64x64x65536_sm75_float16_cm_0_1` was built for SM75, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +} +} +} +} + diff --git a/natten/csrc/autogen/src/cuda/hopper_fmha/source_0.cu b/natten/csrc/autogen/src/cuda/hopper_fmha/source_0.cu new file mode 100644 index 00000000..5f63fc48 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/hopper_fmha/source_0.cu @@ -0,0 +1,222 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_hopper { + + + + +void hopper_fmha_float16_64x128x32( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fmha_hopper::KernelForward< + cutlass::half_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent, cutlass::fmha::collective::DefaultFusion, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::half_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent, cutlass::fmha::collective::ResidualFusion, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::half_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent, cutlass::fmha::collective::CausalFusion, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::half_t, + GemmShape, + natten::cuda::hopper::HopperKernelSchedule::NonPersistent, + cutlass::fmha::collective::ResidualFusion, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::half_t, + GemmShape, + natten::cuda::hopper::HopperKernelSchedule::NonPersistent, + cutlass::fmha::collective::CausalFusion, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void hopper_fmha_float16_64x128x64( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fmha_hopper::KernelForward< + cutlass::half_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent, cutlass::fmha::collective::DefaultFusion, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::half_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent, cutlass::fmha::collective::ResidualFusion, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::half_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent, cutlass::fmha::collective::CausalFusion, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::half_t, + GemmShape, + natten::cuda::hopper::HopperKernelSchedule::NonPersistent, + cutlass::fmha::collective::ResidualFusion, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::half_t, + GemmShape, + natten::cuda::hopper::HopperKernelSchedule::NonPersistent, + cutlass::fmha::collective::CausalFusion, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + +} // namespace fmha_hopper +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/hopper_fmha/source_1.cu b/natten/csrc/autogen/src/cuda/hopper_fmha/source_1.cu new file mode 100644 index 00000000..3090c069 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/hopper_fmha/source_1.cu @@ -0,0 +1,222 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_hopper { + + + + +void hopper_fmha_float16_128x128x128_coop( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fmha_hopper::KernelForward< + cutlass::half_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative, cutlass::fmha::collective::DefaultFusion, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::half_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative, cutlass::fmha::collective::ResidualFusion, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::half_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative, cutlass::fmha::collective::CausalFusion, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::half_t, + GemmShape, + natten::cuda::hopper::HopperKernelSchedule::WSCooperative, + cutlass::fmha::collective::ResidualFusion, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::half_t, + GemmShape, + natten::cuda::hopper::HopperKernelSchedule::WSCooperative, + cutlass::fmha::collective::CausalFusion, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void hopper_fmha_float16_128x128x128_pp( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fmha_hopper::KernelForward< + cutlass::half_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong, cutlass::fmha::collective::DefaultFusion, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::half_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong, cutlass::fmha::collective::ResidualFusion, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::half_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong, cutlass::fmha::collective::CausalFusion, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::half_t, + GemmShape, + natten::cuda::hopper::HopperKernelSchedule::WSPingpong, + cutlass::fmha::collective::ResidualFusion, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::half_t, + GemmShape, + natten::cuda::hopper::HopperKernelSchedule::WSPingpong, + cutlass::fmha::collective::CausalFusion, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + +} // namespace fmha_hopper +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/hopper_fmha/source_2.cu b/natten/csrc/autogen/src/cuda/hopper_fmha/source_2.cu new file mode 100644 index 00000000..29cf191b --- /dev/null +++ b/natten/csrc/autogen/src/cuda/hopper_fmha/source_2.cu @@ -0,0 +1,222 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_hopper { + + + + +void hopper_fmha_float16_128x64x256_coop( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fmha_hopper::KernelForward< + cutlass::half_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative, cutlass::fmha::collective::DefaultFusion, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::half_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative, cutlass::fmha::collective::ResidualFusion, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::half_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative, cutlass::fmha::collective::CausalFusion, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::half_t, + GemmShape, + natten::cuda::hopper::HopperKernelSchedule::WSCooperative, + cutlass::fmha::collective::ResidualFusion, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::half_t, + GemmShape, + natten::cuda::hopper::HopperKernelSchedule::WSCooperative, + cutlass::fmha::collective::CausalFusion, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void hopper_fmha_bfloat16_64x128x32( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fmha_hopper::KernelForward< + cutlass::bfloat16_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent, cutlass::fmha::collective::DefaultFusion, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::bfloat16_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent, cutlass::fmha::collective::ResidualFusion, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::bfloat16_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent, cutlass::fmha::collective::CausalFusion, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::bfloat16_t, + GemmShape, + natten::cuda::hopper::HopperKernelSchedule::NonPersistent, + cutlass::fmha::collective::ResidualFusion, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::bfloat16_t, + GemmShape, + natten::cuda::hopper::HopperKernelSchedule::NonPersistent, + cutlass::fmha::collective::CausalFusion, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + +} // namespace fmha_hopper +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/hopper_fmha/source_3.cu b/natten/csrc/autogen/src/cuda/hopper_fmha/source_3.cu new file mode 100644 index 00000000..30fb7073 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/hopper_fmha/source_3.cu @@ -0,0 +1,222 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_hopper { + + + + +void hopper_fmha_bfloat16_64x128x64( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fmha_hopper::KernelForward< + cutlass::bfloat16_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent, cutlass::fmha::collective::DefaultFusion, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::bfloat16_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent, cutlass::fmha::collective::ResidualFusion, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::bfloat16_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent, cutlass::fmha::collective::CausalFusion, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::bfloat16_t, + GemmShape, + natten::cuda::hopper::HopperKernelSchedule::NonPersistent, + cutlass::fmha::collective::ResidualFusion, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::bfloat16_t, + GemmShape, + natten::cuda::hopper::HopperKernelSchedule::NonPersistent, + cutlass::fmha::collective::CausalFusion, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void hopper_fmha_bfloat16_128x128x128_coop( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fmha_hopper::KernelForward< + cutlass::bfloat16_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative, cutlass::fmha::collective::DefaultFusion, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::bfloat16_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative, cutlass::fmha::collective::ResidualFusion, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::bfloat16_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative, cutlass::fmha::collective::CausalFusion, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::bfloat16_t, + GemmShape, + natten::cuda::hopper::HopperKernelSchedule::WSCooperative, + cutlass::fmha::collective::ResidualFusion, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::bfloat16_t, + GemmShape, + natten::cuda::hopper::HopperKernelSchedule::WSCooperative, + cutlass::fmha::collective::CausalFusion, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + +} // namespace fmha_hopper +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/hopper_fmha/source_4.cu b/natten/csrc/autogen/src/cuda/hopper_fmha/source_4.cu new file mode 100644 index 00000000..8100e10d --- /dev/null +++ b/natten/csrc/autogen/src/cuda/hopper_fmha/source_4.cu @@ -0,0 +1,222 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_hopper { + + + + +void hopper_fmha_bfloat16_128x128x128_pp( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fmha_hopper::KernelForward< + cutlass::bfloat16_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong, cutlass::fmha::collective::DefaultFusion, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::bfloat16_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong, cutlass::fmha::collective::ResidualFusion, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::bfloat16_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong, cutlass::fmha::collective::CausalFusion, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::bfloat16_t, + GemmShape, + natten::cuda::hopper::HopperKernelSchedule::WSPingpong, + cutlass::fmha::collective::ResidualFusion, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::bfloat16_t, + GemmShape, + natten::cuda::hopper::HopperKernelSchedule::WSPingpong, + cutlass::fmha::collective::CausalFusion, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void hopper_fmha_bfloat16_128x64x256_coop( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fmha_hopper::KernelForward< + cutlass::bfloat16_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative, cutlass::fmha::collective::DefaultFusion, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::bfloat16_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative, cutlass::fmha::collective::ResidualFusion, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::bfloat16_t, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative, cutlass::fmha::collective::CausalFusion, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::bfloat16_t, + GemmShape, + natten::cuda::hopper::HopperKernelSchedule::WSCooperative, + cutlass::fmha::collective::ResidualFusion, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_hopper::KernelForward< + cutlass::bfloat16_t, + GemmShape, + natten::cuda::hopper::HopperKernelSchedule::WSCooperative, + cutlass::fmha::collective::CausalFusion, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + +} // namespace fmha_hopper +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/hopper_fmha_bwd/source_0.cu b/natten/csrc/autogen/src/cuda/hopper_fmha_bwd/source_0.cu new file mode 100644 index 00000000..8157d996 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/hopper_fmha_bwd/source_0.cu @@ -0,0 +1,234 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_hopper { + + + + +void hopper_fmha_backward_float16_64x128x32( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fmha_hopper::KernelBackward< + cutlass::half_t, GemmShape, cutlass::fmha::collective::DefaultFusion, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::half_t, GemmShape, cutlass::fmha::collective::ResidualFusion, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::half_t, GemmShape, cutlass::fmha::collective::CausalFusion, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::half_t, + GemmShape, + cutlass::fmha::collective::ResidualFusion, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::half_t, + GemmShape, + cutlass::fmha::collective::CausalFusion, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_q % get<0>(GemmShape{}) == 0 && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void hopper_fmha_backward_float16_128x128x32( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fmha_hopper::KernelBackward< + cutlass::half_t, GemmShape, cutlass::fmha::collective::DefaultFusion, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::half_t, GemmShape, cutlass::fmha::collective::ResidualFusion, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::half_t, GemmShape, cutlass::fmha::collective::CausalFusion, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::half_t, + GemmShape, + cutlass::fmha::collective::ResidualFusion, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::half_t, + GemmShape, + cutlass::fmha::collective::CausalFusion, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_q % get<0>(GemmShape{}) == 0 && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + +} // namespace fmha_hopper +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/hopper_fmha_bwd/source_1.cu b/natten/csrc/autogen/src/cuda/hopper_fmha_bwd/source_1.cu new file mode 100644 index 00000000..d23a62c8 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/hopper_fmha_bwd/source_1.cu @@ -0,0 +1,234 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_hopper { + + + + +void hopper_fmha_backward_float16_64x128x64( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fmha_hopper::KernelBackward< + cutlass::half_t, GemmShape, cutlass::fmha::collective::DefaultFusion, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::half_t, GemmShape, cutlass::fmha::collective::ResidualFusion, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::half_t, GemmShape, cutlass::fmha::collective::CausalFusion, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::half_t, + GemmShape, + cutlass::fmha::collective::ResidualFusion, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::half_t, + GemmShape, + cutlass::fmha::collective::CausalFusion, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_q % get<0>(GemmShape{}) == 0 && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void hopper_fmha_backward_float16_128x128x64( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fmha_hopper::KernelBackward< + cutlass::half_t, GemmShape, cutlass::fmha::collective::DefaultFusion, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::half_t, GemmShape, cutlass::fmha::collective::ResidualFusion, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::half_t, GemmShape, cutlass::fmha::collective::CausalFusion, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::half_t, + GemmShape, + cutlass::fmha::collective::ResidualFusion, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::half_t, + GemmShape, + cutlass::fmha::collective::CausalFusion, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_q % get<0>(GemmShape{}) == 0 && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + +} // namespace fmha_hopper +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/hopper_fmha_bwd/source_2.cu b/natten/csrc/autogen/src/cuda/hopper_fmha_bwd/source_2.cu new file mode 100644 index 00000000..79e469ad --- /dev/null +++ b/natten/csrc/autogen/src/cuda/hopper_fmha_bwd/source_2.cu @@ -0,0 +1,234 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_hopper { + + + + +void hopper_fmha_backward_float16_64x128x128( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fmha_hopper::KernelBackward< + cutlass::half_t, GemmShape, cutlass::fmha::collective::DefaultFusion, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::half_t, GemmShape, cutlass::fmha::collective::ResidualFusion, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::half_t, GemmShape, cutlass::fmha::collective::CausalFusion, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::half_t, + GemmShape, + cutlass::fmha::collective::ResidualFusion, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::half_t, + GemmShape, + cutlass::fmha::collective::CausalFusion, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_q % get<0>(GemmShape{}) == 0 && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void hopper_fmha_backward_bfloat16_64x128x32( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fmha_hopper::KernelBackward< + cutlass::bfloat16_t, GemmShape, cutlass::fmha::collective::DefaultFusion, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::bfloat16_t, GemmShape, cutlass::fmha::collective::ResidualFusion, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::bfloat16_t, GemmShape, cutlass::fmha::collective::CausalFusion, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::bfloat16_t, + GemmShape, + cutlass::fmha::collective::ResidualFusion, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::bfloat16_t, + GemmShape, + cutlass::fmha::collective::CausalFusion, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_q % get<0>(GemmShape{}) == 0 && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + +} // namespace fmha_hopper +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/hopper_fmha_bwd/source_3.cu b/natten/csrc/autogen/src/cuda/hopper_fmha_bwd/source_3.cu new file mode 100644 index 00000000..9248b583 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/hopper_fmha_bwd/source_3.cu @@ -0,0 +1,234 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_hopper { + + + + +void hopper_fmha_backward_bfloat16_128x128x32( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fmha_hopper::KernelBackward< + cutlass::bfloat16_t, GemmShape, cutlass::fmha::collective::DefaultFusion, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::bfloat16_t, GemmShape, cutlass::fmha::collective::ResidualFusion, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::bfloat16_t, GemmShape, cutlass::fmha::collective::CausalFusion, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::bfloat16_t, + GemmShape, + cutlass::fmha::collective::ResidualFusion, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::bfloat16_t, + GemmShape, + cutlass::fmha::collective::CausalFusion, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_q % get<0>(GemmShape{}) == 0 && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void hopper_fmha_backward_bfloat16_64x128x64( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fmha_hopper::KernelBackward< + cutlass::bfloat16_t, GemmShape, cutlass::fmha::collective::DefaultFusion, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::bfloat16_t, GemmShape, cutlass::fmha::collective::ResidualFusion, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::bfloat16_t, GemmShape, cutlass::fmha::collective::CausalFusion, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::bfloat16_t, + GemmShape, + cutlass::fmha::collective::ResidualFusion, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::bfloat16_t, + GemmShape, + cutlass::fmha::collective::CausalFusion, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_q % get<0>(GemmShape{}) == 0 && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + +} // namespace fmha_hopper +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/hopper_fmha_bwd/source_4.cu b/natten/csrc/autogen/src/cuda/hopper_fmha_bwd/source_4.cu new file mode 100644 index 00000000..c16a5e7f --- /dev/null +++ b/natten/csrc/autogen/src/cuda/hopper_fmha_bwd/source_4.cu @@ -0,0 +1,234 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fmha_hopper { + + + + +void hopper_fmha_backward_bfloat16_128x128x64( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fmha_hopper::KernelBackward< + cutlass::bfloat16_t, GemmShape, cutlass::fmha::collective::DefaultFusion, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::bfloat16_t, GemmShape, cutlass::fmha::collective::ResidualFusion, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::bfloat16_t, GemmShape, cutlass::fmha::collective::CausalFusion, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::bfloat16_t, + GemmShape, + cutlass::fmha::collective::ResidualFusion, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::bfloat16_t, + GemmShape, + cutlass::fmha::collective::CausalFusion, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_q % get<0>(GemmShape{}) == 0 && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + + + + +void hopper_fmha_backward_bfloat16_64x128x128( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fmha_hopper::KernelBackward< + cutlass::bfloat16_t, GemmShape, cutlass::fmha::collective::DefaultFusion, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::bfloat16_t, GemmShape, cutlass::fmha::collective::ResidualFusion, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::bfloat16_t, GemmShape, cutlass::fmha::collective::CausalFusion, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::bfloat16_t, + GemmShape, + cutlass::fmha::collective::ResidualFusion, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_hopper::KernelBackward< + cutlass::bfloat16_t, + GemmShape, + cutlass::fmha::collective::CausalFusion, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) { + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_q % get<0>(GemmShape{}) == 0 && seqlen_k % get<1>(GemmShape{}) == 0; + if (no_mask_required) { + Kernel kernel; + launch_kernel(kernel); + } + else if (is_varlen && is_causal) { + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + } + else if (is_varlen && not is_causal) { + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + } + else if (not is_varlen && is_causal) { + KernelWithCausalMask kernel; + launch_kernel(kernel); + } + else { + KernelWithResidualMask kernel; + launch_kernel(kernel); + } +} + + +} // namespace fmha_hopper +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/hopper_fna/source_0.cu b/natten/csrc/autogen/src/cuda/hopper_fna/source_0.cu new file mode 100644 index 00000000..14b0f488 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/hopper_fna/source_0.cu @@ -0,0 +1,1882 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_hopper { + + + + +void hopper_fna1d_float16_64x128x32_Q64_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_float16_64x128x32_Q64_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_float16_64x128x64_Q64_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_float16_64x128x64_Q64_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_float16_128x128x128_coop_Q128_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_float16_128x128x128_pp_Q128_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_float16_128x128x128_coop_Q128_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_float16_128x128x128_pp_Q128_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_float16_128x64x256_coop_Q128_KV64_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_float16_128x64x256_coop_Q128_KV64_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_bfloat16_64x128x32_Q64_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_bfloat16_64x128x32_Q64_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_bfloat16_64x128x64_Q64_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_bfloat16_64x128x64_Q64_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_bfloat16_128x128x128_coop_Q128_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_bfloat16_128x128x128_pp_Q128_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_bfloat16_128x128x128_coop_Q128_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_bfloat16_128x128x128_pp_Q128_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_bfloat16_128x64x256_coop_Q128_KV64_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_bfloat16_128x64x256_coop_Q128_KV64_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_64x128x32_Q8x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_64x128x32_Q8x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_64x128x32_Q8x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_64x128x32_Q8x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_64x128x32_Q8x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_64x128x32_Q8x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_64x128x32_Q8x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_64x128x32_Q8x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_64x128x64_Q8x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_64x128x64_Q8x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_64x128x64_Q8x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_hopper +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/hopper_fna/source_1.cu b/natten/csrc/autogen/src/cuda/hopper_fna/source_1.cu new file mode 100644 index 00000000..22609680 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/hopper_fna/source_1.cu @@ -0,0 +1,1882 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_hopper { + + + + +void hopper_fna2d_float16_64x128x64_Q8x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_64x128x64_Q8x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_64x128x64_Q8x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_64x128x64_Q8x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_64x128x64_Q8x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_128x128x128_coop_Q16x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_128x128x128_pp_Q16x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_128x128x128_coop_Q16x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_128x128x128_pp_Q16x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_128x128x128_coop_Q16x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_128x128x128_pp_Q16x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_128x128x128_coop_Q16x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_128x128x128_pp_Q16x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_128x64x256_coop_Q16x8_KV8x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_128x64x256_coop_Q8x16_KV8x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_128x64x256_coop_Q16x8_KV8x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_128x64x256_coop_Q8x16_KV8x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_128x64x256_coop_Q16x8_KV8x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_128x64x256_coop_Q8x16_KV8x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_128x64x256_coop_Q16x8_KV8x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_float16_128x64x256_coop_Q8x16_KV8x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_64x128x32_Q8x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_64x128x32_Q8x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_64x128x32_Q8x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_64x128x32_Q8x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_64x128x32_Q8x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_64x128x32_Q8x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_64x128x32_Q8x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_64x128x32_Q8x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_64x128x64_Q8x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_64x128x64_Q8x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_hopper +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/hopper_fna/source_2.cu b/natten/csrc/autogen/src/cuda/hopper_fna/source_2.cu new file mode 100644 index 00000000..2d9dc6e7 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/hopper_fna/source_2.cu @@ -0,0 +1,1882 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_hopper { + + + + +void hopper_fna2d_bfloat16_64x128x64_Q8x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_64x128x64_Q8x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_64x128x64_Q8x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_64x128x64_Q8x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_64x128x64_Q8x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_64x128x64_Q8x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_128x128x128_coop_Q16x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_128x128x128_pp_Q16x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_128x128x128_coop_Q16x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_128x128x128_pp_Q16x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_128x128x128_coop_Q16x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_128x128x128_pp_Q16x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_128x128x128_coop_Q16x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_128x128x128_pp_Q16x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_128x64x256_coop_Q16x8_KV8x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_128x64x256_coop_Q8x16_KV8x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_128x64x256_coop_Q16x8_KV8x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_128x64x256_coop_Q8x16_KV8x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_128x64x256_coop_Q16x8_KV8x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_128x64x256_coop_Q8x16_KV8x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_128x64x256_coop_Q16x8_KV8x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_bfloat16_128x64x256_coop_Q8x16_KV8x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<16>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_hopper +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/hopper_fna/source_3.cu b/natten/csrc/autogen/src/cuda/hopper_fna/source_3.cu new file mode 100644 index 00000000..5be53167 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/hopper_fna/source_3.cu @@ -0,0 +1,1882 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_hopper { + + + + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x32_Q4x4x4_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_64x128x64_Q4x4x4_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_coop_Q4x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_pp_Q4x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_coop_Q2x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_pp_Q2x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_coop_Q4x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_pp_Q4x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_coop_Q2x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_pp_Q2x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_hopper +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/hopper_fna/source_4.cu b/natten/csrc/autogen/src/cuda/hopper_fna/source_4.cu new file mode 100644 index 00000000..5129172d --- /dev/null +++ b/natten/csrc/autogen/src/cuda/hopper_fna/source_4.cu @@ -0,0 +1,1822 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_hopper { + + + + +void hopper_fna3d_float16_128x128x128_coop_Q4x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_pp_Q4x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_coop_Q2x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_pp_Q2x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_coop_Q4x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_pp_Q4x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_coop_Q2x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_pp_Q2x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_coop_Q4x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_pp_Q4x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_coop_Q2x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_pp_Q2x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_coop_Q4x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_pp_Q4x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_coop_Q2x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_pp_Q2x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_coop_Q4x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_pp_Q4x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_coop_Q2x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_pp_Q2x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_coop_Q4x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_pp_Q4x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_coop_Q2x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x128x128_pp_Q2x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x64x256_coop_Q4x4x8_KV4x4x4_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x64x256_coop_Q2x8x8_KV4x4x4_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x64x256_coop_Q4x4x8_KV4x4x4_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x64x256_coop_Q2x8x8_KV4x4x4_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x64x256_coop_Q4x4x8_KV4x4x4_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x64x256_coop_Q2x8x8_KV4x4x4_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_hopper +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/hopper_fna/source_5.cu b/natten/csrc/autogen/src/cuda/hopper_fna/source_5.cu new file mode 100644 index 00000000..429b83d6 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/hopper_fna/source_5.cu @@ -0,0 +1,1822 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_hopper { + + + + +void hopper_fna3d_float16_128x64x256_coop_Q4x4x8_KV4x4x4_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x64x256_coop_Q2x8x8_KV4x4x4_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x64x256_coop_Q4x4x8_KV4x4x4_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x64x256_coop_Q2x8x8_KV4x4x4_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x64x256_coop_Q4x4x8_KV4x4x4_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x64x256_coop_Q2x8x8_KV4x4x4_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x64x256_coop_Q4x4x8_KV4x4x4_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x64x256_coop_Q2x8x8_KV4x4x4_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x64x256_coop_Q4x4x8_KV4x4x4_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_float16_128x64x256_coop_Q2x8x8_KV4x4x4_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_hopper +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/hopper_fna/source_6.cu b/natten/csrc/autogen/src/cuda/hopper_fna/source_6.cu new file mode 100644 index 00000000..c3dd6fba --- /dev/null +++ b/natten/csrc/autogen/src/cuda/hopper_fna/source_6.cu @@ -0,0 +1,1822 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_hopper { + + + + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::NonPersistent>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_coop_Q4x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_pp_Q4x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_coop_Q2x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_pp_Q2x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_coop_Q4x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_pp_Q4x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_coop_Q2x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_pp_Q2x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_coop_Q4x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_pp_Q4x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_coop_Q2x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_pp_Q2x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_coop_Q4x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_pp_Q4x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_coop_Q2x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_pp_Q2x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_coop_Q4x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_pp_Q4x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_hopper +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/hopper_fna/source_7.cu b/natten/csrc/autogen/src/cuda/hopper_fna/source_7.cu new file mode 100644 index 00000000..40623b59 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/hopper_fna/source_7.cu @@ -0,0 +1,1822 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_hopper { + + + + +void hopper_fna3d_bfloat16_128x128x128_coop_Q2x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_pp_Q2x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_coop_Q4x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_pp_Q4x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_coop_Q2x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_pp_Q2x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_coop_Q4x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_pp_Q4x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_coop_Q2x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_pp_Q2x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_coop_Q4x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_pp_Q4x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_coop_Q2x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x128x128_pp_Q2x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSPingpong>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x64x256_coop_Q4x4x8_KV4x4x4_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x64x256_coop_Q2x8x8_KV4x4x4_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x64x256_coop_Q4x4x8_KV4x4x4_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x64x256_coop_Q2x8x8_KV4x4x4_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x64x256_coop_Q4x4x8_KV4x4x4_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x64x256_coop_Q2x8x8_KV4x4x4_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x64x256_coop_Q4x4x8_KV4x4x4_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x64x256_coop_Q2x8x8_KV4x4x4_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x64x256_coop_Q4x4x8_KV4x4x4_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x64x256_coop_Q2x8x8_KV4x4x4_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x64x256_coop_Q4x4x8_KV4x4x4_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x64x256_coop_Q2x8x8_KV4x4x4_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x64x256_coop_Q4x4x8_KV4x4x4_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x64x256_coop_Q2x8x8_KV4x4x4_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x64x256_coop_Q4x4x8_KV4x4x4_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_bfloat16_128x64x256_coop_Q2x8x8_KV4x4x4_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using GemmShape = cute::tuple, cute::Int<64>, cute::Int<256>>; + using Kernel = natten::cuda::fna_hopper::KernelForward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape, natten::cuda::hopper::HopperKernelSchedule::WSCooperative>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_hopper +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/hopper_fna_bwd/source_0.cu b/natten/csrc/autogen/src/cuda/hopper_fna_bwd/source_0.cu new file mode 100644 index 00000000..f10b1666 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/hopper_fna_bwd/source_0.cu @@ -0,0 +1,4986 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_hopper { + + + + +void hopper_fna1d_backward_float16_64x128x32_Q64_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_backward_float16_128x128x32_Q128_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_backward_float16_64x128x32_Q64_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_backward_float16_128x128x32_Q128_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_backward_float16_64x128x64_Q64_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_backward_float16_128x128x64_Q128_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_backward_float16_64x128x64_Q64_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_backward_float16_128x128x64_Q128_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_backward_float16_64x128x128_Q64_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_backward_float16_64x128x128_Q64_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_backward_bfloat16_64x128x32_Q64_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_backward_bfloat16_128x128x32_Q128_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_backward_bfloat16_64x128x32_Q64_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_backward_bfloat16_128x128x32_Q128_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_backward_bfloat16_64x128x64_Q64_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_backward_bfloat16_128x128x64_Q128_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_backward_bfloat16_64x128x64_Q64_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_backward_bfloat16_128x128x64_Q128_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_backward_bfloat16_64x128x128_Q64_KV128_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna1d_backward_bfloat16_64x128x128_Q64_KV128_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple>; + using KVTileShape = cute::tuple>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_64x128x32_Q8x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_64x128x32_Q8x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_128x128x32_Q16x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_128x128x32_Q16x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_64x128x32_Q8x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_64x128x32_Q8x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_128x128x32_Q16x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_128x128x32_Q16x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_64x128x32_Q8x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_64x128x32_Q8x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_128x128x32_Q16x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_128x128x32_Q16x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_64x128x32_Q8x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_64x128x32_Q8x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_128x128x32_Q16x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_128x128x32_Q16x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_64x128x64_Q8x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_64x128x64_Q8x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_128x128x64_Q16x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_128x128x64_Q16x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_64x128x64_Q8x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_64x128x64_Q8x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_128x128x64_Q16x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_128x128x64_Q16x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_64x128x64_Q8x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_64x128x64_Q8x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_128x128x64_Q16x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_128x128x64_Q16x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_64x128x64_Q8x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_64x128x64_Q8x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_128x128x64_Q16x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_128x128x64_Q16x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_64x128x128_Q8x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_64x128x128_Q8x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_64x128x128_Q8x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_64x128x128_Q8x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_64x128x128_Q8x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_64x128x128_Q8x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_64x128x128_Q8x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_float16_64x128x128_Q8x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_64x128x32_Q8x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_64x128x32_Q8x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_128x128x32_Q16x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_128x128x32_Q16x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_64x128x32_Q8x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_64x128x32_Q8x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_128x128x32_Q16x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_128x128x32_Q16x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_64x128x32_Q8x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_64x128x32_Q8x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_128x128x32_Q16x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_128x128x32_Q16x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_64x128x32_Q8x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_hopper +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/hopper_fna_bwd/source_1.cu b/natten/csrc/autogen/src/cuda/hopper_fna_bwd/source_1.cu new file mode 100644 index 00000000..1069d8e3 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/hopper_fna_bwd/source_1.cu @@ -0,0 +1,4986 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_hopper { + + + + +void hopper_fna2d_backward_bfloat16_64x128x32_Q8x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_128x128x32_Q16x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_128x128x32_Q16x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_64x128x64_Q8x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_64x128x64_Q8x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_128x128x64_Q16x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_128x128x64_Q16x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_64x128x64_Q8x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_64x128x64_Q8x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_128x128x64_Q16x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_128x128x64_Q16x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_64x128x64_Q8x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_64x128x64_Q8x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_128x128x64_Q16x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_128x128x64_Q16x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_64x128x64_Q8x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_64x128x64_Q8x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_128x128x64_Q16x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_128x128x64_Q16x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_64x128x128_Q8x8_KV16x8_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_64x128x128_Q8x8_KV8x16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_64x128x128_Q8x8_KV16x8_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_64x128x128_Q8x8_KV8x16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_64x128x128_Q8x8_KV16x8_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_64x128x128_Q8x8_KV8x16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_64x128x128_Q8x8_KV16x8_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna2d_backward_bfloat16_64x128x128_Q8x8_KV8x16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<16>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x32_Q4x4x4_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x32_Q4x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_hopper +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/hopper_fna_bwd/source_2.cu b/natten/csrc/autogen/src/cuda/hopper_fna_bwd/source_2.cu new file mode 100644 index 00000000..582fa5a8 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/hopper_fna_bwd/source_2.cu @@ -0,0 +1,4986 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_hopper { + + + + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x64_Q4x4x4_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_128x128x64_Q4x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q2x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q1x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q2x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q1x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q2x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q1x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q2x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q1x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q2x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q1x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q2x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q1x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q2x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q1x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q4x4x4_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q2x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_float16_64x128x128_Q1x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::half_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_hopper +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/hopper_fna_bwd/source_3.cu b/natten/csrc/autogen/src/cuda/hopper_fna_bwd/source_3.cu new file mode 100644 index 00000000..47a62761 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/hopper_fna_bwd/source_3.cu @@ -0,0 +1,4986 @@ +#ifdef NATTEN_WITH_CUTLASS +#ifdef NATTEN_WITH_HOPPER_FNA +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace fna_hopper { + + + + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x32_Q4x4x4_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x32_Q4x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<32>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x64_Q4x4x4_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_128x128x64_Q4x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<64>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV4x4x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q2x4x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q1x8x8_KV2x8x8_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV4x4x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q2x4x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q1x8x8_KV2x8x8_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV4x4x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q2x4x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q1x8x8_KV2x8x8_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV4x4x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q2x4x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q1x8x8_KV2x8x8_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV4x4x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q2x4x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q1x8x8_KV2x8x8_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV4x4x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q2x4x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q1x8x8_KV2x8x8_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV4x4x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q2x4x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q1x8x8_KV2x8x8_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV4x4x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q4x4x4_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<4>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q2x4x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<4>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + + + + +void hopper_fna3d_backward_bfloat16_64x128x128_Q1x8x8_KV2x8x8_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + cute::tuple q_shape, + cute::tuple kv_shape, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) { + + using Causal = cute::tuple; + using QTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using KVTileShape = cute::tuple, cute::Int<8>, cute::Int<8>>; + using GemmShape = cute::tuple, cute::Int<128>, cute::Int<128>>; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + cutlass::bfloat16_t, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({bytes}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +} + + +} // namespace fna_hopper +} // namespace cuda +} // namespace natten +#endif +#endif + diff --git a/natten/csrc/autogen/src/cuda/reference/source_0.cu b/natten/csrc/autogen/src/cuda/reference/source_0.cu new file mode 100644 index 00000000..36797114 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/reference/source_0.cu @@ -0,0 +1,2181 @@ +#ifdef NATTEN_WITH_CUTLASS +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace reference { + + + + +void reference_fna1d_float32_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna1d_float32_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna1d_float16_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna1d_float16_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna1d_bfloat16_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna1d_bfloat16_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna1d_backward_float32_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna1d_backward_float32_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna1d_backward_float16_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna1d_backward_float16_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna1d_backward_bfloat16_causal0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna1d_backward_bfloat16_causal1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna2d_float32_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna2d_float32_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna2d_float32_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna2d_float32_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna2d_float16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna2d_float16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna2d_float16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna2d_float16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna2d_bfloat16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna2d_bfloat16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna2d_bfloat16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna2d_bfloat16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna2d_backward_float32_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna2d_backward_float32_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna2d_backward_float32_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna2d_backward_float32_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna2d_backward_float16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna2d_backward_float16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna2d_backward_float16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna2d_backward_float16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna2d_backward_bfloat16_causal0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna2d_backward_bfloat16_causal0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna2d_backward_bfloat16_causal1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna2d_backward_bfloat16_causal1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_float32_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_float32_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_float32_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_float32_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_float32_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_float32_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + +} // namespace reference +} // namespace cuda +} // namespace natten +#endif + diff --git a/natten/csrc/autogen/src/cuda/reference/source_1.cu b/natten/csrc/autogen/src/cuda/reference/source_1.cu new file mode 100644 index 00000000..e36b6a78 --- /dev/null +++ b/natten/csrc/autogen/src/cuda/reference/source_1.cu @@ -0,0 +1,2229 @@ +#ifdef NATTEN_WITH_CUTLASS +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace natten { +namespace cuda { +namespace reference { + + + + +void reference_fna3d_float32_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_float32_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_float16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_float16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_float16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_float16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_float16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_float16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_float16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_float16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_bfloat16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_bfloat16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_bfloat16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_bfloat16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_bfloat16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_bfloat16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_bfloat16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_bfloat16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_forward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_backward_float32_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_backward_float32_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_backward_float32_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_backward_float32_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_backward_float32_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_backward_float32_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_backward_float32_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_backward_float32_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_backward_float16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_backward_float16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_backward_float16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_backward_float16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_backward_float16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_backward_float16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_backward_float16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_backward_float16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_backward_bfloat16_causal0x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_backward_bfloat16_causal0x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_backward_bfloat16_causal0x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_backward_bfloat16_causal0x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_backward_bfloat16_causal1x0x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_backward_bfloat16_causal1x0x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_backward_bfloat16_causal1x1x0( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + + + + +void reference_fna3d_backward_bfloat16_causal1x1x1( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + cute::tuple qkv_shape, + cute::tuple window_size, + cute::tuple stride, + cute::tuple dilation, + float attn_scale, + cudaStream_t stream) { + + using Causal = cute::tuple; + + fna_reference_backward( + static_cast(ptr_Q), + static_cast(ptr_K), + static_cast(ptr_V), + static_cast(ptr_O), + static_cast(ptr_DO), + static_cast(ptr_DQ), + static_cast(ptr_DK), + static_cast(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{}, + attn_scale, + stream); +} + + +} // namespace reference +} // namespace cuda +} // namespace natten +#endif + diff --git a/natten/csrc/include/natten/blackwell_fmha.h b/natten/csrc/include/natten/blackwell_fmha.h new file mode 100644 index 00000000..a332bf78 --- /dev/null +++ b/natten/csrc/include/natten/blackwell_fmha.h @@ -0,0 +1,74 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ +/*! \file + \brief Blackwell FMHA interface +*/ + +#pragma once +#include + +#include + +namespace natten { + +void blackwell_fmha_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::optional& logsumexp, + bool is_causal, + float attn_scale, + int query_tile_size, + int key_tile_size, + bool run_persistent, + // varlen + const at::optional& cumulative_seqlen_Q, + const at::optional& cumulative_seqlen_KV, + // only used if cumulative_seqlen_Q and cumulative_seqlen_KV are specified + int max_seqlen_Q, + int max_seqlen_KV); + +void blackwell_fmha_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + bool is_causal, + float attn_scale, + int query_tile_size, + int key_tile_size, + // varlen + const at::optional& cumulative_seqlen_Q, + const at::optional& cumulative_seqlen_KV, + // only used if cumulative_seqlen_Q and cumulative_seqlen_KV are specified + int max_seqlen_Q, + int max_seqlen_KV, + bool deterministic); + +} // namespace natten diff --git a/natten/csrc/include/natten/blackwell_fna.h b/natten/csrc/include/natten/blackwell_fna.h new file mode 100644 index 00000000..44118e90 --- /dev/null +++ b/natten/csrc/include/natten/blackwell_fna.h @@ -0,0 +1,155 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ +/*! \file + \brief Blackwell FNA interface +*/ + +#pragma once +#include + +#include + +namespace natten { + +// Forward + +void blackwell_na1d_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::optional& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& q_shape_, + const std::tuple& kv_shape_, + const std::tuple& qkv_shape_, + const std::tuple& query_tile_shape, + const std::tuple& key_tile_shape, + bool run_persistent); + +void blackwell_na2d_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::optional& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& q_shape_, + const std::tuple& kv_shape_, + const std::tuple& qkv_shape_, + const std::tuple& query_tile_shape, + const std::tuple& key_tile_shape, + bool run_persistent); + +void blackwell_na3d_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::optional& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& q_shape_, + const std::tuple& kv_shape_, + const std::tuple& qkv_shape_, + const std::tuple& query_tile_shape, + const std::tuple& key_tile_shape, + bool run_persistent); + +// Backward + +void blackwell_na1d_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& q_shape_, + const std::tuple& kv_shape_, + const std::tuple& qkv_shape_, + const std::tuple& query_tile_shape, + const std::tuple& key_tile_shape); + +void blackwell_na2d_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& q_shape_, + const std::tuple& kv_shape_, + const std::tuple& qkv_shape_, + const std::tuple& query_tile_shape, + const std::tuple& key_tile_shape); + +void blackwell_na3d_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& q_shape_, + const std::tuple& kv_shape_, + const std::tuple& qkv_shape_, + const std::tuple& query_tile_shape, + const std::tuple& key_tile_shape); + +} // namespace natten diff --git a/natten/csrc/include/natten/compute_delta.h b/natten/csrc/include/natten/compute_delta.h new file mode 100644 index 00000000..db17a693 --- /dev/null +++ b/natten/csrc/include/natten/compute_delta.h @@ -0,0 +1,40 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ +/*! \file + \brief compute_delta interface + mostly used to test the kernel. +*/ + +#pragma once +#include + +#include + +namespace natten { + +void compute_delta( + const at::Tensor& out, + const at::Tensor& d_out, + at::Tensor& delta); + +} // namespace natten diff --git a/natten/csrc/include/natten/cuda/fmha/epilogue/epilogue_pipelined.h b/natten/csrc/include/natten/cuda/fmha/epilogue/epilogue_pipelined.h new file mode 100644 index 00000000..3b06c109 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha/epilogue/epilogue_pipelined.h @@ -0,0 +1,638 @@ +/* + * Copied from xFormers (https://github.com/facebookresearch/xformers/) + * + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + */ +/*************************************************************************************************** + * Copyright (c) 2017 - 2026 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 Epilogue for threadblock scoped GEMMs using Tensor Ops. + + File copied from "cutlass/epilogue/threadblock/epilogue.h" + then modified to: + (1) load 2 source fragments at the same time (pipelining) + (2) support reading from a different dtype + (3) pass the row id to the OutputOp if it takes it + (see MemoryEfficientAttentionNormalize) + Note that in general the fragment passed to the OutputOp could + span multiple rows but it does not happen with the configurations we have +*/ + +#pragma once + +#if defined(__CUDACC_RTC__) +#include +#else +#include +#endif + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include + +#include +#include +#include + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace epilogue { +namespace threadblock { + +template +struct ApplyEpilogueOp { + static CUTLASS_DEVICE typename Op::FragmentOutput apply( + Op const& output_op, + int row_id, + typename Op::FragmentAccumulator const& accum, + typename Op::FragmentOutput const& source) { + return output_op(accum, source); + } + static CUTLASS_DEVICE typename Op::FragmentOutput apply( + Op const& output_op, + int row_id, + typename Op::FragmentAccumulator const& accum) { + return output_op(accum); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Epilogue operator +template < + typename Shape_, ///< Shape of threadblock tile (concept: GemmShape) + typename WarpMmaOperator_, ///< Warp-level MMA operator (concept: + ///< gemm::warp::MmaTensorOp) + int PartitionsK, ///< Number of partitions of the K dimension + typename OutputTileIterator_, ///< Tile iterator writing output tensors + typename AccumulatorFragmentIterator_, ///< Fragment iterator selecting + ///< accumulators + typename WarpTileIterator_, ///< Warp-scoped tile iterator writing + ///< accumulators to SMEM + typename SharedLoadIterator_, ///< Threadblock-scoped tile iterator loading + ///< from SMEM + typename OutputOp_, ///< Output operator + typename Padding_, ///< Padding added to SMEM allocation to avoid bank + ///< conflicts (concept: MatrixShape) + int FragmentsPerPartition = + 1, ///< Used to coarsten the epilogue granularity + int IterationsUnroll = ///< Used to reduce binary size when epilogue op is + ///< large + (!IsEpilogueFunctorHeavy::value), + typename OutputTileSourceIterator_ = + OutputTileIterator_ ///< Tile iterator reading tensors + > +class EpiloguePipelined : public EpilogueBase< + Shape_, + typename WarpMmaOperator_::Shape, + PartitionsK, + AccumulatorFragmentIterator_, + WarpTileIterator_, + Padding_, + FragmentsPerPartition> { + public: + using Base = EpilogueBase< + Shape_, + typename WarpMmaOperator_::Shape, + PartitionsK, + AccumulatorFragmentIterator_, + WarpTileIterator_, + Padding_, + FragmentsPerPartition>; + + using Shape = Shape_; + using WarpMmaOperator = WarpMmaOperator_; + static int const kPartitionsK = PartitionsK; + using OutputTileIterator = OutputTileIterator_; + using OutputTileSourceIterator = OutputTileSourceIterator_; + using AccumulatorFragmentIterator = AccumulatorFragmentIterator_; + using WarpTileIterator = WarpTileIterator_; + using SharedLoadIterator = SharedLoadIterator_; + using OutputOp = OutputOp_; + using Padding = Padding_; + + using Layout = layout::RowMajor; + using LongIndex = typename Layout::LongIndex; + + /// The complete warp-level accumulator tile + using AccumulatorTile = typename Base::AccumulatorTile; + + /// Accumulator element + using ElementAccumulator = typename WarpTileIterator::Element; + + /// Output element + using ElementOutput = typename OutputTileIterator::Element; + using ElementSource = typename OutputTileSourceIterator::Element; + + /// Output access size + static int const kElementsPerAccess = OutputTileIterator::kElementsPerAccess; + + /// Tensor reference to destination tensor + using TensorRef = typename OutputTileIterator::TensorRef; + + /// Tensor reference to sync tensor + using SyncTensorRef = + typename cutlass::TensorRef; + + /// Const tensor reference to source tensor + using ConstTensorRef = typename OutputTileIterator::ConstTensorRef; + + /// Array type used to output + using OutputAccessType = Array< + typename OutputTileIterator::Element, + OutputTileIterator::kElementsPerAccess>; + using SourceAccessType = Array< + typename OutputTileSourceIterator::Element, + OutputTileSourceIterator::kElementsPerAccess>; + + /// Array type used by output functor + using AccumulatorAccessType = Array< + typename WarpTileIterator::Element, + OutputTileIterator::kElementsPerAccess>; + + /// Number of warps + using WarpCount = typename Base::WarpCount; + + static int constexpr kSmemTiles = Base::kFragmentsPerIteration > 1 + ? Base::kFragmentsPerIteration + : kPartitionsK; + static int constexpr kSmemPointerOffset = + Base::SharedStorage::StorageShape::kCount / kSmemTiles; + + public: + static_assert( + OutputTileSourceIterator::Fragment::kElements == + OutputTileIterator::Fragment::kElements, + "Mismatch between input tile and output tile iterator (kElements)"); + static_assert( + OutputTileSourceIterator::kIterations == OutputTileIterator::kIterations, + "Mismatch between input tile and output tile iterator (kIterations)"); + static_assert( + SharedLoadIterator::Fragment::kElements == + OutputTileIterator::Fragment::kElements, + "Mismatch between shared load iterator and output tile iterator."); + + static_assert( + OutputTileIterator::kElementsPerAccess, + "OutputTileIterator::kElementsPerAccess must not be zero."); + + static_assert( + !(OutputTileIterator::Fragment::kElements % + OutputTileIterator::kElementsPerAccess), + "Divisibility"); + + private: + /// Loads fragment from shared memory aligned with output tensor + SharedLoadIterator shared_load_iterator_; + + public: + /// Constructor + CUTLASS_DEVICE + EpiloguePipelined( + typename Base::SharedStorage& shared_storage, ///< Shared storage object + int thread_idx, ///< ID of a thread within the threadblock + int warp_idx, ///< ID of warp within threadblock + int lane_idx ///< Id of thread within warp + ) + : Base(shared_storage, thread_idx, warp_idx, lane_idx), + shared_load_iterator_(shared_storage.reference(), thread_idx) {} + + /// Streams the result to global memory + CUTLASS_DEVICE + void operator()( + OutputOp const& output_op, ///< Output operator + OutputTileIterator + destination_iterator, ///< Tile iterator for destination + AccumulatorTile const& + accumulators, ///< Complete warp-level accumulator tile + OutputTileSourceIterator + source_iterator) { ///< Threadblock tile coordinate in GEMM (in units + ///< of threadblock tiles) + + if (!output_op.is_source_needed()) { + compute_source_not_needed_(output_op, destination_iterator, accumulators); + } else { + compute_source_needed_( + output_op, destination_iterator, accumulators, source_iterator); + } + } + CUTLASS_DEVICE + void operator()( + OutputOp const& output_op, ///< Output operator + OutputTileIterator + destination_iterator, ///< Tile iterator for destination + AccumulatorTile const& + accumulators) { ///< Complete warp-level accumulator tile + compute_source_not_needed_(output_op, destination_iterator, accumulators); + } + + private: + template + struct acc2smem_source_not_needed; + + template + struct acc2smem_source_not_needed> { + template + CUTLASS_DEVICE static void helper( + AccumulatorFragmentIterator accum_fragment_iterator, + WarpTileIterator& warp_tile_iterator) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < Advance; i++) { + ++accum_fragment_iterator; + } + + CUTLASS_PRAGMA_UNROLL + for (int p = 0; p < Base::kFragmentsPerIteration; ++p) { + typename AccumulatorFragmentIterator::Fragment accum_fragment; + + accum_fragment_iterator.load(accum_fragment); + ++accum_fragment_iterator; + + warp_tile_iterator.store(accum_fragment); + if (p < Base::kFragmentsPerIteration - 1) { + warp_tile_iterator.add_pointer_offset(kSmemPointerOffset); + } + } + + if (Base::kFragmentsPerIteration > 1) { + warp_tile_iterator.add_pointer_offset( + kSmemPointerOffset * (1 - Base::kFragmentsPerIteration)); + } + } + + CUTLASS_DEVICE + static void push( + size_t pos, + AccumulatorFragmentIterator const& iterator_begin, + WarpTileIterator& warp_tile_iterator) { + int dummy[] = { + (pos == (Seq * Base::kFragmentsPerIteration)) && + (helper( + iterator_begin, warp_tile_iterator), + 0)...}; + + CUTLASS_UNUSED(dummy[0]); + } + }; + + static_assert( + kPartitionsK == 1 || Base::kFragmentsPerIteration == 1, + "One of these must be exactly 1."); + + /// Streams the result to global memory + CUTLASS_DEVICE + void compute_source_not_needed_( + OutputOp const& output_op, ///< Output operator + OutputTileIterator + destination_iterator, ///< Tile iterator for destination + AccumulatorTile const& + accumulators ///< Complete warp-level accumulator tile + ) { + // + // Iterator over warp-level accumulator fragment + // + + AccumulatorFragmentIterator accum_fragment_iterator(accumulators); + + // + // Iterate over accumulator tile + // + +#pragma unroll( \ + IterationsUnroll \ + ? OutputTileIterator::kIterations / Base::kFragmentsPerIteration \ + : 1) + for (int iter = 0; iter < OutputTileIterator::kIterations; + iter += Base::kFragmentsPerIteration) { + // + // Convert and store fragment + // + + __syncthreads(); + + acc2smem_source_not_needed>:: + push(iter, accum_fragment_iterator, this->warp_tile_iterator_); + + __syncthreads(); + + // + // Load fragments from shared memory + // + + CUTLASS_PRAGMA_UNROLL + for (int p = 0; p < Base::kFragmentsPerIteration; ++p) { + typename SharedLoadIterator::Fragment + aligned_accum_fragment[kPartitionsK]; + + shared_load_iterator_.load(aligned_accum_fragment[0]); + + if (p < Base::kFragmentsPerIteration - 1) { + shared_load_iterator_.add_pointer_offset(kSmemPointerOffset); + } else if (kPartitionsK > 1) { + plus add_fragments; + + CUTLASS_PRAGMA_UNROLL + for (int i = 1; i < kPartitionsK; ++i) { + shared_load_iterator_.add_pointer_offset(kSmemPointerOffset); + shared_load_iterator_.load(aligned_accum_fragment[i]); + aligned_accum_fragment[0] = add_fragments( + aligned_accum_fragment[0], aligned_accum_fragment[i]); + } + + shared_load_iterator_.add_pointer_offset( + (1 - kPartitionsK) * kSmemPointerOffset); + } + + // + // Compute the output result + // + + typename OutputTileIterator::Fragment output_fragment; + + apply_output_operator_source_not_needed_( + destination_iterator.thread_start_row(), + output_fragment, + output_op, + aligned_accum_fragment[0]); + + // + // Store the final result + // + + destination_iterator.store(output_fragment); + ++destination_iterator; + } + + if (Base::kFragmentsPerIteration > 1) { + shared_load_iterator_.add_pointer_offset( + kSmemPointerOffset * (1 - Base::kFragmentsPerIteration)); + } + } + } + + template + struct acc2smem_source_needed; + + template + struct acc2smem_source_needed> { + template + CUTLASS_DEVICE static void helper( + AccumulatorFragmentIterator accum_fragment_iterator, + WarpTileIterator& warp_tile_iterator) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < Advance; i++) { + ++accum_fragment_iterator; + } + + typename AccumulatorFragmentIterator::Fragment accum_fragment; + accum_fragment_iterator.load(accum_fragment); + warp_tile_iterator.store(accum_fragment); + } + + CUTLASS_DEVICE + static void push( + size_t pos, + AccumulatorFragmentIterator const& iterator_begin, + WarpTileIterator& warp_tile_iterator) { + int dummy[] = { + (pos == Seq) && + (helper(iterator_begin, warp_tile_iterator), 0)...}; + } + }; + + /// Streams the result to global memory + CUTLASS_DEVICE + void compute_source_needed_( + OutputOp const& output_op, ///< Output operator + OutputTileIterator + destination_iterator, ///< Tile iterator for destination + AccumulatorTile const& + accumulators, ///< Complete warp-level accumulator tile + OutputTileSourceIterator + source_iterator ///< Threadblock tile coordinate in GEMM (in units of + ///< threadblock tiles) + ) { + typename OutputTileSourceIterator::Fragment source_fragment[2]; + + source_fragment[0].clear(); + source_iterator.load(source_fragment[0]); + ++source_iterator; + source_fragment[1].clear(); + + // + // Iterator over warp-level accumulator fragment + // + + AccumulatorFragmentIterator accum_fragment_iterator(accumulators); + + // + // Iterate over accumulator tile + // + +#pragma unroll(IterationsUnroll ? OutputTileIterator::kIterations : 1) + for (int iter = 0; iter < OutputTileIterator::kIterations; ++iter) { + if (iter > 0) { + __syncthreads(); + } + // + // Load the source for next iteration (pipelining) + // + + if (iter + 1 < OutputTileIterator::kIterations) { + source_iterator.load(source_fragment[(iter + 1) % 2]); + } + ++source_iterator; + acc2smem_source_needed< + cutlass::make_index_sequence>:: + push(iter, accum_fragment_iterator, this->warp_tile_iterator_); + + __syncthreads(); + + // + // Load fragments from shared memory + // + + typename SharedLoadIterator::Fragment + aligned_accum_fragment[kPartitionsK]; + + shared_load_iterator_.load(aligned_accum_fragment[0]); + + // If the number of k-slices is > 1 - perform a reduction amongst the + // k-slices + if (kPartitionsK > 1) { + plus add_fragments; + + CUTLASS_PRAGMA_UNROLL + for (int i = 1; i < kPartitionsK; ++i) { + shared_load_iterator_.add_pointer_offset(kSmemPointerOffset); + shared_load_iterator_.load(aligned_accum_fragment[i]); + aligned_accum_fragment[0] = add_fragments( + aligned_accum_fragment[0], aligned_accum_fragment[i]); + } + + shared_load_iterator_.add_pointer_offset( + (1 - kPartitionsK) * kSmemPointerOffset); + } + + // + // Compute the output result + // + + typename OutputTileIterator::Fragment output_fragment; + + apply_output_operator_( + destination_iterator.thread_start_row(), + output_fragment, + output_op, + aligned_accum_fragment[0], + source_fragment[iter % 2]); + + // + // Store the final result + // + + destination_iterator.store(output_fragment); + ++destination_iterator; + } + } + + /// Helper to invoke the output functor over each vector of output + CUTLASS_DEVICE + void apply_output_operator_( + int begin_row, + typename OutputTileIterator::Fragment& output_fragment, + OutputOp const& output_op, ///< Output operator + typename SharedLoadIterator::Fragment const& aligned_accum_fragment, + typename OutputTileSourceIterator::Fragment const& source_fragment) { + OutputAccessType* output_frag_ptr = + reinterpret_cast(&output_fragment); + + AccumulatorAccessType const* compute_frag_ptr = + reinterpret_cast(&aligned_accum_fragment); + + SourceAccessType const* source_frag_ptr = + reinterpret_cast(&source_fragment); + + int const kOutputOpIterations = OutputTileIterator::Fragment::kElements / + OutputTileIterator::kElementsPerAccess; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kOutputOpIterations; ++i) { + // Call the output operator + output_frag_ptr[i] = ApplyEpilogueOp::apply( + output_op, + begin_row + getRowOffset(i * OutputTileIterator::kElementsPerAccess), + compute_frag_ptr[i], + source_frag_ptr[i]); + } + } + + /// Helper to invoke the output functor over each vector of output + CUTLASS_DEVICE + void apply_output_operator_source_not_needed_( + int begin_row, + typename OutputTileIterator::Fragment& output_fragment, + OutputOp const& output_op, ///< Output operator + typename SharedLoadIterator::Fragment const& aligned_accum_fragment) { + OutputAccessType* output_frag_ptr = + reinterpret_cast(&output_fragment); + + AccumulatorAccessType const* compute_frag_ptr = + reinterpret_cast(&aligned_accum_fragment); + + int const kOutputOpIterations = OutputTileIterator::Fragment::kElements / + OutputTileIterator::kElementsPerAccess; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kOutputOpIterations; ++i) { + // Call the output operator + output_frag_ptr[i] = ApplyEpilogueOp::apply( + output_op, + begin_row + getRowOffset(i * OutputTileIterator::kElementsPerAccess), + compute_frag_ptr[i]); + } + } + + // This should be constexpr, but it's only supported on c++14 + static int CUTLASS_HOST_DEVICE getRowOffset(int i) { + using ThreadMap = typename OutputTileIterator::ThreadMap; + + CUTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; + ++cluster) { + CUTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + CUTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + int row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + int frag_row_idx = + (row + + ThreadMap::Iterations::kRow * + (group + ThreadMap::Iterations::kGroup * cluster)); + CUTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; + ++column) { + int frag_idx = ThreadMap::kElementsPerAccess * + (frag_row_idx * ThreadMap::Iterations::kColumn + column); + if (i < frag_idx + ThreadMap::kElementsPerAccess) { + return row_offset; + } + } + } + } + } + return -1; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fmha/epilogue/epilogue_rescale_output.h b/natten/csrc/include/natten/cuda/fmha/epilogue/epilogue_rescale_output.h new file mode 100644 index 00000000..8f9fdf6a --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha/epilogue/epilogue_rescale_output.h @@ -0,0 +1,267 @@ +/* + * Copied from xFormers (https://github.com/facebookresearch/xformers/) + * + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * BSD 3-Clause License + * + * 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 names of Facebook, Deepmind Technologies, NYU, NEC + * Laboratories America and IDIAP Research Institute 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 OWNER 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 Epilogue for threadblock scoped GEMMs using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory + to match canonical tensor layouts in global memory. Epilogues support + conversion and reduction operations. + + This is a copy of cutlass/epilogue/threadblock/epilogue.h that can + handle "row_id" as a first argument, as uses it to get the corresponding + `m_prime` / `s_prime` to rescale the output. +*/ + +#pragma once + +#if defined(__CUDACC_RTC__) +#include +#else +#include +#endif + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include + +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace epilogue { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Applies a linear combination operator to an array of elements. +// output <- alpha * accumulator + beta * source +// with: +// alpha = 1 / s_prime (to normalize when isLast=True, 1 otherwise) +// beta = alpha / m_prime (renormalize the output when the max changes) +// source is the current output +template < + typename ElementOutput_, ///< Data type used to store tensors + typename ElementSource_, //< Data type for source (usually matches + //`ElementOutput`) + int Count, ///< Number of elements computed per operation. + ///< Usually it is 128/sizeof_bits, + ///< but we use 64 or 32 sometimes when there are not enough data + ///< to store + typename ElementAccumulator_, ///< Accumulator data type + typename ElementCompute_, ///< Data type used to compute linear combination + bool isFirst, + bool isLast, + typename FragmentAlphaBeta_, + FloatRoundStyle Round = FloatRoundStyle::round_to_nearest> +class MemoryEfficientAttentionNormalize { + public: + using ElementOutput = ElementOutput_; + using ElementSource = ElementSource_; + using ElementAccumulator = ElementAccumulator_; + using ElementCompute = ElementCompute_; + + static int const kCount = Count; + + using FragmentOutput = Array; + using FragmentSource = Array; + using FragmentAccumulator = Array; + using ComputeFragment = Array; + using FragmentAlphaBeta = FragmentAlphaBeta_; + + static FloatRoundStyle const kRound = Round; + + private: + // + // Data members + // + + FragmentAlphaBeta const& s_prime_; + FragmentAlphaBeta const& m_prime_; + + public: + /// Constructs the function object, possibly loading from pointers in host + /// memory + CUTLASS_HOST_DEVICE + MemoryEfficientAttentionNormalize( + FragmentAlphaBeta const& s_prime, + FragmentAlphaBeta const& m_prime) + : s_prime_(s_prime), m_prime_(m_prime) {} + + /// Returns true if source is needed + CUTLASS_HOST_DEVICE + bool is_source_needed() const { + return !isFirst; + } + + /// Functionally required for serial reduction in the epilogue + CUTLASS_HOST_DEVICE + void set_k_partition(int k_partition, int k_partition_count) {} + + /// Computes linear scaling: D = alpha * accumulator + beta * source + CUTLASS_HOST_DEVICE + FragmentOutput operator()( + int row, + FragmentAccumulator const& accumulator, + FragmentSource const& source) const { + assert(!isFirst); + + // Convert source to interal compute numeric type + NumericArrayConverter + source_converter; + NumericArrayConverter + accumulator_converter; + + // Convert to destination numeric type + NumericArrayConverter + destination_converter; + + ComputeFragment converted_source = source_converter(source); + ComputeFragment converted_accumulator = accumulator_converter(accumulator); + + // Perform binary operations + ComputeFragment intermediate; + + multiplies mul_add_source; + multiply_add mul_add_accumulator; + + ElementCompute alpha = isLast ? (1 / s_prime_[row]) : 1; + ElementCompute beta = alpha * m_prime_[row]; + + intermediate = mul_add_source(beta, converted_source); // X = beta * C + + intermediate = mul_add_accumulator( + alpha, converted_accumulator, intermediate); // D = alpha * Accum + X + + return destination_converter(intermediate); + } + + /// Computes linear scaling: D = alpha * accumulator + CUTLASS_HOST_DEVICE + FragmentOutput operator()(int row, FragmentAccumulator const& accumulator) + const { + assert(isFirst); + + // Convert source to interal compute numeric type + NumericArrayConverter + accumulator_converter; + + // Convert to destination numeric type + NumericArrayConverter + destination_converter; + + ComputeFragment converted_accumulator = accumulator_converter(accumulator); + + ComputeFragment intermediate; + multiplies mul_accumulator; + + ElementCompute alpha = isLast ? (1 / s_prime_[row]) : 1; + + intermediate = mul_accumulator( + alpha, converted_accumulator); // X = alpha * C + uniform + + return destination_converter(intermediate); + } +}; + +} // namespace thread + +namespace threadblock { +template < + typename EO, + typename ES, + int Count, + typename EA, + typename EC, + bool F, + bool L, + typename FAB, + FloatRoundStyle R> +struct ApplyEpilogueOp> { + using Op = thread:: + MemoryEfficientAttentionNormalize; + static CUTLASS_DEVICE typename Op::FragmentOutput apply( + Op const& output_op, + int row_id, + typename Op::FragmentAccumulator const& accum, + typename Op::FragmentSource const& source) { + return output_op(row_id, accum, source); + } + static CUTLASS_DEVICE typename Op::FragmentOutput apply( + Op const& output_op, + int row_id, + typename Op::FragmentAccumulator const& accum) { + return output_op(row_id, accum); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fmha/epilogue/epilogue_thread_apply_logsumexp.h b/natten/csrc/include/natten/cuda/fmha/epilogue/epilogue_thread_apply_logsumexp.h new file mode 100644 index 00000000..725327ba --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha/epilogue/epilogue_thread_apply_logsumexp.h @@ -0,0 +1,183 @@ +/* + * Copied from xFormers (https://github.com/facebookresearch/xformers/) + * + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + */ +/*************************************************************************************************** + * Copyright (c) 2017 - 2026 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 Functor performing linear combination operations used by epilogues. +*/ + +#pragma once + +#include + +#include "cutlass/array.h" +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/thread/activation.h" +#include "cutlass/functional.h" +#include "cutlass/numeric_conversion.h" +#include "cutlass/numeric_types.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace epilogue { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace detail { + +template +struct ArrayExponential { + CUTLASS_HOST_DEVICE + Array operator()( + Array const& input) const { + Array result; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < ElementsPerAccess; ++i) { + result[i] = expf(input[i]); + } + + return result; + } +}; + +template +struct ArrayExponential { + CUTLASS_DEVICE + Array operator()( + Array const& input) const { + Array result; + + int const kVectorCount = ElementsPerAccess / 2; + + __half2 const* input_ptr = + reinterpret_cast<__half2 const*>(input.raw_data()); + __half2* res_ptr = reinterpret_cast<__half2*>(result.raw_data()); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kVectorCount; ++i) { + res_ptr[i] = h2exp(input_ptr[i]); + } + + return result; + } +}; +} // namespace detail + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Applies: +/// output <- (input).exp() +template < + typename ElementOutput_, // output + typename ElementLSE_, // accumulator from LSE + typename ElementAccumulator_, // accumulator from matmul + typename ElementCompute_, // intermediate compute (and exp calculation) + int ElementsPerAccess> +// class ApplyLogSumExp { +class ApplyExp { + public: + using ElementOutput = ElementOutput_; + using ElementAccumulator = ElementAccumulator_; + using ElementCompute = ElementCompute_; + using ElementLSE = ElementLSE_; + + static int const kElementsPerAccess = ElementsPerAccess; + static int const kCount = kElementsPerAccess; + static const ScaleType::Kind kScale = + cutlass::epilogue::thread::ScaleType::NoBetaScaling; + + using FragmentOutput = Array; + using FragmentAccumulator = Array; + using FragmentCompute = Array; + using FragmentLSE = Array; + using FragmentScaleBias = FragmentLSE; // Used by epilogue_smem_accumulator.h + + public: + // + // Methods + // + + CUTLASS_HOST_DEVICE + ApplyExp() {} + + /// Returns true if source is needed + CUTLASS_HOST_DEVICE + bool is_source_needed() const { + return true; + } + + /// Functionally required for serial reduction in the epilogue + CUTLASS_HOST_DEVICE + void set_k_partition(int k_partition, int k_partition_count) {} + + CUTLASS_HOST_DEVICE + FragmentOutput operator()( + FragmentAccumulator const& AB/*, + FragmentLSE const& scale_unused, + // bias used as LSE + FragmentLSE const& bias*/) const { + FragmentCompute frag_AB = NumericArrayConverter< + ElementCompute, + ElementAccumulator, + kElementsPerAccess>()(AB); + // FragmentCompute frag_lse_compute = + // NumericArrayConverter()( + // bias); + FragmentCompute frag_compute; + + // minus minus_lse; + detail::ArrayExponential apply_exp; + // frag_compute = minus_lse(frag_AB, frag_lse_compute); + frag_compute = apply_exp(frag_AB); + + return NumericArrayConverter< + ElementOutput, + ElementCompute, + kElementsPerAccess>()(frag_compute); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace epilogue +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fmha/fmha_backward.cuh b/natten/csrc/include/natten/cuda/fmha/fmha_backward.cuh new file mode 100644 index 00000000..be3d8826 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha/fmha_backward.cuh @@ -0,0 +1,239 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Fused Neighborhood Attention kernels are heavily based on the + * memory-efficient attention kernels from the xFormers project by Meta + * Platforms, Inc. + * + * Copyright (c) Facebook, Inc. and its affiliates + * + * BSD 3-Clause License + * + * 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 names of Facebook, Deepmind Technologies, NYU, NEC + * Laboratories America and IDIAP Research Institute 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 OWNER 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 + +#include + +namespace natten { +namespace cuda { +namespace fmha { + +template +void fmha_backward_generic( + at::ScalarType dtype, + const int cc, + const size_t max_smem, + cudaStream_t stream, + MemoryAllocator alloc_bytes, + void* grad_out_ptr, + void* query_ptr, + void* key_ptr, + void* value_ptr, + void* logsumexp_ptr, + void* delta_ptr, + void* out_ptr, + // Outputs: + void* grad_query_ptr, + void* grad_key_ptr, + void* grad_value_ptr, + // Params + int32_t batch_size, + int32_t seqlen_q, + int32_t seqlen_kv, + int32_t heads, + int32_t dim, + int32_t dim_value, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int query_tile_size, + int key_tile_size, + int num_splits_key) { + bool kernel_launched = false; + const auto maxK = std::max(dim, dim_value); + + auto launchKernel = [&](auto _k, auto kernel_fn) { + using Kernel = decltype(_k); + using scalar_t = typename Kernel::scalar_t; + (void)_k; + + if (kernel_launched) { + return; + } + // Check if this kernel is compatible + if (Kernel::kMaxK < maxK) { + return; + } + + if (Kernel::kKeysQueriesAlignedToBlockSize && + (is_varlen || seqlen_q % Kernel::kBlockSizeI || + seqlen_kv % Kernel::kBlockSizeJ)) { + return; + } + + if (query_tile_size != Kernel::kBlockSizeI || + key_tile_size != Kernel::kBlockSizeJ) { + return; + } + + if ((dim % Kernel::kMinimumAlignment) || + (dim % Kernel::kMinimumAlignment) || + (dim_value % Kernel::kMinimumAlignment)) { + return; + } + // Uses too much shmem + size_t smem_bytes = sizeof(typename Kernel::SharedStorage); + if (smem_bytes > max_smem) { + return; + } + + kernel_launched = true; + + typename Kernel::Params p; + p.query_ptr = (scalar_t*)query_ptr; + p.key_ptr = (scalar_t*)key_ptr; + p.value_ptr = (scalar_t*)value_ptr; + p.logsumexp_ptr = (typename Kernel::lse_scalar_t*)logsumexp_ptr; + p.output_ptr = (scalar_t*)out_ptr; + p.grad_output_ptr = (scalar_t*)grad_out_ptr; + p.grad_query_ptr = (scalar_t*)grad_query_ptr; + p.grad_key_ptr = (scalar_t*)grad_key_ptr; + p.grad_value_ptr = (scalar_t*)grad_value_ptr; + p.delta_ptr = (float*)delta_ptr; + + p.scale = attn_scale; + + p.num_heads = heads; + p.head_dim = dim; + p.head_dim_value = dim_value; + p.num_queries = seqlen_q; + p.num_keys = seqlen_kv; + p.num_batches = batch_size; + p.custom_mask_type = is_causal ? 1 : 0; + + if (is_varlen) { + p.num_queries = max_seqlen_Q; + p.num_keys = max_seqlen_KV; + p.ptr_cumulative_seqlen_Q = + static_cast(ptr_cumulative_seqlen_Q); + p.ptr_cumulative_seqlen_KV = + static_cast(ptr_cumulative_seqlen_KV); + } + + p.num_splits_key = num_splits_key; + + p.o_strideH = dim_value; + p.q_strideH = dim; + p.k_strideH = dim; + p.v_strideH = dim_value; + p.o_strideB = seqlen_q * heads * dim_value; + p.q_strideB = seqlen_q * heads * dim; + p.k_strideB = seqlen_kv * heads * dim; + p.v_strideB = seqlen_kv * heads * dim_value; + p.lse_strideB = seqlen_q * heads; + p.delta_strideB = seqlen_q * heads; + + p.gO_strideB = seqlen_q * heads * dim_value; + p.gQ_strideB = seqlen_q * heads * dim; + p.gK_strideB = seqlen_kv * heads * dim; + p.gV_strideB = seqlen_kv * heads * dim_value; + p.gO_strideH = dim_value; + p.gQ_strideH = dim; + p.gK_strideH = dim; + p.gV_strideH = dim_value; + p.q_strideM = heads * dim; + p.k_strideM = heads * dim; + p.v_strideM = heads * dim_value; + p.gO_strideM = heads * dim_value; + + int64_t size_bytes = p.workspace_size(); + if (size_bytes) { + void* workspace_ptr = nullptr; + alloc_bytes( + &workspace_ptr, size_bytes, true /*p.should_zero_workspace()*/); + p.workspace = (float*)workspace_ptr; + } + + Kernel::check_supported(p); + + if (smem_bytes > 0xc000) { + // https://docs.nvidia.com/cuda/cuda-c-programming-guide/#features-and-technical-specifications-technical-specifications-per-compute-capability + auto err = cudaFuncSetAttribute( + kernel_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes); + NATTEN_CHECK( + err != cudaErrorInvalidValue, + "This GPU does not have enough shared-memory."); + NATTEN_CUDA_CHECK(err); + } + + // second syntax resulted in the error below on windows + // error C3495: 'kernel_fn': a simple capture must be a variable + // with automatic storage duration declared + // in the reaching scope of the lambda +#ifdef NATTEN_WINDOWS + cudaFuncAttributes attr; + NATTEN_CUDA_CHECK(cudaFuncGetAttributes(&attr, kernel_fn)); + NATTEN_CHECK( + attr.binaryVersion >= Kernel::ArchTag::kMinComputeCapability, + "Something went wrong in the build process"); +#else + auto checkBinaryArchMatches = [&]() { + cudaFuncAttributes attr; + NATTEN_CUDA_CHECK(cudaFuncGetAttributes(&attr, kernel_fn)); + return attr.binaryVersion >= Kernel::ArchTag::kMinComputeCapability; + }; + NATTEN_CHECK( + checkBinaryArchMatches(), "Something went wrong in the build process"); +#endif + + kernel_fn<<>>(p); + }; + + DISPATCH_FMHA_BACKWARD(cc, dtype, launchKernel); + NATTEN_CHECK( + kernel_launched, "Could not find a compatible FMHA backward kernel."); + NATTEN_CUDA_CHECK(cudaGetLastError()); +} + +} // namespace fmha +} // namespace cuda +} // namespace natten diff --git a/natten/csrc/include/natten/cuda/fmha/fmha_forward.cuh b/natten/csrc/include/natten/cuda/fmha/fmha_forward.cuh new file mode 100644 index 00000000..34c36d78 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha/fmha_forward.cuh @@ -0,0 +1,191 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Fused Neighborhood Attention kernels are heavily based on the + * memory-efficient attention kernels from the xFormers project by Meta + * Platforms, Inc. + * + * Copyright (c) Facebook, Inc. and its affiliates + * + * BSD 3-Clause License + * + * 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 names of Facebook, Deepmind Technologies, NYU, NEC + * Laboratories America and IDIAP Research Institute 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 OWNER 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 + +#include + +namespace natten { +namespace cuda { +namespace fmha { + +template +void fmha_forward_generic( + at::ScalarType dtype, + const int cc, + const size_t max_smem, + cudaStream_t stream, + MemoryAllocator alloc_bytes, + void* query_ptr, + void* key_ptr, + void* value_ptr, + void* out_ptr, + int32_t batch_size, + int32_t seqlen_q, + int32_t seqlen_kv, + int32_t heads, + int32_t dim, + int32_t dim_value, + bool is_causal, + float attn_scale, + void* logsumexp_ptr, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int query_tile_size, + int key_tile_size) { + bool kernel_launched = false; + auto launchKernel = [&](auto _k, auto kernel_fn) { + using Kernel = decltype(_k); + using scalar_t = typename Kernel::scalar_t; + (void)_k; + + if (kernel_launched) { + return; + } + + if (dim_value > Kernel::kMaxK || dim > Kernel::kMaxK) { + return; + } + + if (query_tile_size != Kernel::kQueriesPerBlock || + key_tile_size != Kernel::kKeysPerBlock) { + return; + } + + // Alignment + if ((dim % Kernel::kAlignmentQ) || (dim % Kernel::kAlignmentK) || + (dim_value % Kernel::kAlignmentV)) { + return; + } + // Uses too much smem + size_t smem_bytes = sizeof(typename Kernel::SharedStorage); + if (smem_bytes > max_smem) { + return; + } + + kernel_launched = true; + + typename Kernel::Params p; + p.query_ptr = (scalar_t*)query_ptr; + p.key_ptr = (scalar_t*)key_ptr; + p.value_ptr = (scalar_t*)value_ptr; + p.logsumexp_ptr = logsumexp_ptr != nullptr + ? (typename Kernel::lse_scalar_t*)logsumexp_ptr + : nullptr; + + p.output_ptr = (typename Kernel::output_t*)out_ptr; + + p.num_heads = heads; + p.head_dim = dim; + p.head_dim_value = dim_value; + p.num_queries = seqlen_q; + p.num_keys = seqlen_kv; + p.num_batches = batch_size; + p.custom_mask_type = is_causal ? 1 : 0; + + if (is_varlen) { + p.num_queries = max_seqlen_Q; + p.num_keys = max_seqlen_KV; + p.ptr_cumulative_seqlen_Q = + static_cast(ptr_cumulative_seqlen_Q); + p.ptr_cumulative_seqlen_KV = + static_cast(ptr_cumulative_seqlen_KV); + } + + if (Kernel::kNeedsOutputAccumulatorBuffer) { + using AccumType = typename Kernel::output_accum_t; + int64_t workspace_size_bytes = static_cast(batch_size) * + static_cast(p.num_queries) * static_cast(heads) * + static_cast(dim_value) * + static_cast(sizeof(AccumType)); + void* accum_ptr = nullptr; + alloc_bytes(&accum_ptr, workspace_size_bytes, /* zero_fill = */ false); + p.output_accum_ptr = (AccumType*)accum_ptr; + } else { + p.output_accum_ptr = nullptr; + } + + p.scale = attn_scale; + + p.q_strideM = heads * dim; + p.k_strideM = heads * dim; + p.v_strideM = heads * dim_value; + p.o_strideM = heads * dim_value; + + p.q_strideH = dim; + p.k_strideH = dim; + p.v_strideH = dim_value; + + p.q_strideB = seqlen_q * heads * dim; + p.k_strideB = seqlen_kv * heads * dim; + p.v_strideB = seqlen_kv * heads * dim_value; + + if (smem_bytes > 0xc000) { + auto err = cudaFuncSetAttribute( + kernel_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes); + NATTEN_CHECK( + err != cudaErrorInvalidValue, + "This GPU does not have enough shared-memory."); + NATTEN_CUDA_CHECK(err); + } + auto blocks = p.getBlocksGrid(); + Kernel::check_supported(p); + kernel_fn<<>>(p); + }; + + DISPATCH_FMHA_FORWARD(cc, dtype, launchKernel); + NATTEN_CHECK( + kernel_launched, "Could not find a compatible FMHA forward kernel."); + NATTEN_CUDA_CHECK(cudaGetLastError()); +} + +} // namespace fmha +} // namespace cuda +} // namespace natten diff --git a/natten/csrc/include/natten/cuda/fmha/gemm/custom_mma.h b/natten/csrc/include/natten/cuda/fmha/gemm/custom_mma.h new file mode 100644 index 00000000..6d89795a --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha/gemm/custom_mma.h @@ -0,0 +1,131 @@ +/* + * Copied from xFormers (https://github.com/facebookresearch/xformers/) and + * edited + * + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * BSD 3-Clause License + * + * 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 names of Facebook, Deepmind Technologies, NYU, NEC + * Laboratories America and IDIAP Research Institute 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 OWNER 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 +#include + +template +struct MakeCustomMma; + +template < + typename Shape, + typename IteratorA, + typename SmemIteratorA, + cutlass::arch::CacheOperation::Kind CacheOpA, + typename IteratorB, + typename SmemIteratorB, + cutlass::arch::CacheOperation::Kind CacheOpB, + typename ElementC, + typename LayoutC, + typename Policy, + int Stages, + cutlass::gemm::SharedMemoryClearOption SharedMemoryClear, + int kMaxK> +struct MakeCustomMma< + cutlass::gemm::threadblock::MmaMultistage< + Shape, + IteratorA, + SmemIteratorA, + CacheOpA, + IteratorB, + SmemIteratorB, + CacheOpB, + ElementC, + LayoutC, + Policy, + Stages, + SharedMemoryClear>, + kMaxK> { + // Reduce the number of stages if we don't need that many + static int constexpr kStages = + kMaxK == cutlass::platform::numeric_limits::max() + ? Stages + : cutlass::const_min( + Stages, + (kMaxK + int(Shape::kK) - 1) / int(Shape::kK)); + using Mma = cutlass::gemm::threadblock::CustomMmaMultistage< + Shape, + IteratorA, + SmemIteratorA, + CacheOpA, + IteratorB, + SmemIteratorB, + CacheOpB, + ElementC, + LayoutC, + Policy, + kStages, + SharedMemoryClear, + kMaxK>; +}; + +template < + typename Shape, + typename IteratorA, + typename SmemIteratorA, + typename IteratorB, + typename SmemIteratorB, + typename ElementC, + typename LayoutC, + typename Policy, + int kMaxK> +struct MakeCustomMma< + cutlass::gemm::threadblock::MmaPipelined< + Shape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + Policy>, + kMaxK> { + using Mma = cutlass::gemm::threadblock::CustomMmaPipelined< + Shape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + Policy>; +}; diff --git a/natten/csrc/include/natten/cuda/fmha/gemm/custom_mma_base.h b/natten/csrc/include/natten/cuda/fmha/gemm/custom_mma_base.h new file mode 100644 index 00000000..3da8befc --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha/gemm/custom_mma_base.h @@ -0,0 +1,189 @@ +/* + * Copied from xFormers (https://github.com/facebookresearch/xformers/) + * + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + */ +/*************************************************************************************************** + * Copyright (c) 2017 - 2026 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 Template for a double-buffered threadblock-scoped GEMM kernel. +*/ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace gemm { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Used for partial specialization + typename Enable = bool> +class CustomMmaBase { + public: + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + ///< Policy describing tuning details + using Policy = Policy_; + + // + // Dependent types + // + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Shape describing the overall GEMM computed from shared memory + /// by each warp. + using WarpGemm = typename Policy::Operator::Shape; + + /// Shape describing the number of warps filling the CTA + using WarpCount = GemmShape< + Shape::kM / WarpGemm::kM, + Shape::kN / WarpGemm::kN, + Shape::kK / WarpGemm::kK>; + + /// Number of warp-level GEMM oeprations + static int const kWarpGemmIterations = + (WarpGemm::kK / Operator::Policy::MmaShape::kK); + + /// Number of stages + static int const kStages = Stages; + + // + // Nested structs + // + + /// Shared storage object needed by threadblock-scoped GEMM + template + struct OperandSharedStorage { + AlignedBuffer buffer; + using TensorRef = TensorRef; + + CUTLASS_DEVICE + static OperandLayout Layout() { + return OperandLayout::packed({OperandShape::kRow, OperandShape::kColumn}); + } + + /// Returns a TensorRef to the operand + CUTLASS_HOST_DEVICE + TensorRef ref() { + return TensorRef{buffer.data(), Layout()}; + } + }; + + /// Shape of the A matrix operand in shared memory + using ShapeA = MatrixShape< + Shape::kM + Policy::SmemPaddingA::kRow, + Shape::kK * kStages + Policy::SmemPaddingA::kColumn>; + + /// Shape of the B matrix operand in shared memory + using ShapeB = MatrixShape< + Shape::kK * kStages + Policy::SmemPaddingB::kRow, + Shape::kN + Policy::SmemPaddingB::kColumn>; + + using SharedStorageA = OperandSharedStorage< + typename Operator::ElementA, + ShapeA, + typename Operator::LayoutA>; + using SharedStorageB = OperandSharedStorage< + typename Operator::ElementB, + ShapeB, + typename Operator::LayoutB>; + using TensorRefA = typename SharedStorageA::TensorRef; + using TensorRefB = typename SharedStorageB::TensorRef; + + struct SharedStorage { + /// Buffer for A operand + SharedStorageA operand_A; + + /// Buffer for B operand + SharedStorageB operand_B; + }; + + protected: + // + // Data members + // + + /// Iterator to load a warp-scoped tile of A operand from shared memory + typename Operator::IteratorA warp_tile_iterator_A_; + + /// Iterator to load a warp-scoped tile of B operand from shared memory + typename Operator::IteratorB warp_tile_iterator_B_; + + public: + /// Construct from tensor references + CUTLASS_DEVICE + CustomMmaBase( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + SharedStorageA& shared_storageA, + SharedStorageB& shared_storageB, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx) + : warp_tile_iterator_A_(shared_storageA.ref(), lane_idx), + warp_tile_iterator_B_(shared_storageB.ref(), lane_idx) {} +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fmha/gemm/custom_mma_multistage.h b/natten/csrc/include/natten/cuda/fmha/gemm/custom_mma_multistage.h new file mode 100644 index 00000000..46c91b6e --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha/gemm/custom_mma_multistage.h @@ -0,0 +1,765 @@ +/* + * Copied from xFormers (https://github.com/facebookresearch/xformers/) + * + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + */ +/*************************************************************************************************** + * Copyright (c) 2017 - 2026 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 Template for a double-buffered threadblock-scoped GEMM kernel. +*/ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace gemm { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + /// Cache operation for operand A + cutlass::arch::CacheOperation::Kind CacheOpA, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Cache operation for operand B + cutlass::arch::CacheOperation::Kind CacheOpB, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone, + /// Upper boundon the K dimension + int kMaxK = cutlass::platform::numeric_limits::max(), + /// Used for partial specialization + typename Enable = bool> +class CustomMmaMultistage : public CustomMmaBase { + public: + ///< Base class + using Base = CustomMmaBase; + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + ///< Iterates over tiles of A operand in global memory + using IteratorA = IteratorA_; + ///< Iterates over tiles of B operand in global memory + using IteratorB = IteratorB_; + ///< Data type of accumulator matrix + using ElementC = ElementC_; + ///< Layout of accumulator matrix + using LayoutC = LayoutC_; + ///< Policy describing tuning details + using Policy = Policy_; + + using SmemIteratorA = SmemIteratorA_; + using SmemIteratorB = SmemIteratorB_; + + static cutlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA; + static cutlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB; + + // + // Dependent types + // + + /// Fragment of accumulator tile + using FragmentC = typename Policy::Operator::FragmentC; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Minimum architecture is Sm80 to support cp.async + using ArchTag = arch::Sm80; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = Operator::kTransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = Operator::kTransformB; + + /// Internal structure exposed for introspection. + struct Detail { + static_assert( + Base::kWarpGemmIterations > 1, + "The pipelined structure requires at least two warp-level " + "GEMM operations."); + + /// Number of cp.async instructions to load one stage of operand A + static int const AsyncCopyIterationsPerStageA = + IteratorA::ThreadMap::Iterations::kCount; + + /// Number of cp.async instructions to load one stage of operand B + static int const AsyncCopyIterationsPerStageB = + IteratorB::ThreadMap::Iterations::kCount; + + /// Number of stages + static int const kStages = Stages; + + /// Number of cp.async instructions to load on group of operand A + static int const kAccessesPerGroupA = + (AsyncCopyIterationsPerStageA + Base::kWarpGemmIterations - 1) / + Base::kWarpGemmIterations; + + /// Number of cp.async instructions to load on group of operand B + static int const kAccessesPerGroupB = + (AsyncCopyIterationsPerStageB + Base::kWarpGemmIterations - 1) / + Base::kWarpGemmIterations; + }; + + static bool const kSmemContainsEntireMat = kMaxK <= Shape::kK * Stages; + static constexpr int kNumStagesConcurrentLoad = + kSmemContainsEntireMat ? Stages : Stages - 1; + + private: + using WarpLoadedFragmentA = typename Operator::FragmentA; + using WarpLoadedFragmentB = typename Operator::FragmentB; + using WarpTransformedFragmentA = typename Operator::TransformedFragmentA; + using WarpTransformedFragmentB = typename Operator::TransformedFragmentB; + + private: + // + // Data members + // + + /// Iterator to write threadblock-scoped tile of A operand to shared memory + SmemIteratorA smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB smem_iterator_B_; + + bool prologue_done_; + + // Set to `True` to ensure the accumulator will be zero outside the GEMM + // footprint + bool zero_outside_bounds_; + + public: + /// Construct from tensor references + CUTLASS_DEVICE + CustomMmaMultistage( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + typename Base::SharedStorageA& shared_storageA, + typename Base::SharedStorageB& shared_storageB, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx) + : Base(shared_storageA, shared_storageB, thread_idx, warp_idx, lane_idx), + smem_iterator_A_(shared_storageA.ref(), thread_idx), + smem_iterator_B_(shared_storageB.ref(), thread_idx), + prologue_done_(false), + zero_outside_bounds_(false) { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + + int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; + int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset( + {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); + } + CUTLASS_DEVICE + CustomMmaMultistage( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + typename Base::SharedStorage& st, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx) + : CustomMmaMultistage( + st.operand_A, + st.operand_B, + thread_idx, + warp_idx, + lane_idx) {} + + CUTLASS_DEVICE + void set_prologue_done(bool value) { + prologue_done_ = value; + } + + CUTLASS_DEVICE + void set_zero_outside_bounds(bool value) { + zero_outside_bounds_ = value; + } + + template + CUTLASS_DEVICE static void prologue( + typename Base::SharedStorage& shared_storage, + ///< iterator over A operand in global memory + IteratorA iterator_A, + ///< iterator over B operand in global memory + IteratorB iterator_B, + int thread_idx, + int problem_size_k) { + prologue( + shared_storage.operand_A, + shared_storage.operand_B, + iterator_A, + iterator_B, + thread_idx, + problem_size_k); + } + + template + CUTLASS_DEVICE static void prologue( + typename Base::SharedStorageA& shared_storageA, + typename Base::SharedStorageB& shared_storageB, + ///< iterator over A operand in global memory + IteratorA iterator_A, + ///< iterator over B operand in global memory + IteratorB iterator_B, + int thread_idx, + int problem_size_k) { + SmemIteratorA smem_iterator_A(shared_storageA.ref(), thread_idx); + SmemIteratorB smem_iterator_B(shared_storageB.ref(), thread_idx); + int32_t iter = (problem_size_k + Base::Shape::kK - 1) / Base::Shape::kK; + _prologue( + iterator_A, iterator_B, iter, smem_iterator_A, smem_iterator_B); + } + + CUTLASS_DEVICE + void copy_tiles_and_advance( + IteratorA& iterator_A, + IteratorB& iterator_B, + int group_start_A = 0, + int group_start_B = 0) { + iterator_A.set_iteration_index( + group_start_A * IteratorA::kAccessesPerVector); + this->smem_iterator_A_.set_iteration_index(group_start_A); + + // Async Copy for operand A + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupA; ++j) { + if (group_start_A + j < Detail::AsyncCopyIterationsPerStageA) { + typename IteratorA::AccessType* dst_ptr = + reinterpret_cast( + this->smem_iterator_A_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / + IteratorA::kAccessesPerVector / 8; + + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { + auto gmem_ptr = iterator_A.get(); + + if (zero_outside_bounds_ || + SharedMemoryClear == SharedMemoryClearOption::kZfill) { + cutlass::arch::cp_async_zfill( + dst_ptr + v, gmem_ptr, iterator_A.valid()); + } else { + cutlass::arch::cp_async( + dst_ptr + v, gmem_ptr, iterator_A.valid()); + } + + ++iterator_A; + } + + ++this->smem_iterator_A_; + } + } + + iterator_B.set_iteration_index( + group_start_B * IteratorB::kAccessesPerVector); + this->smem_iterator_B_.set_iteration_index(group_start_B); + + // Async Copy for operand B + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupB; ++j) { + if (group_start_B + j < Detail::AsyncCopyIterationsPerStageB) { + typename IteratorB::AccessType* dst_ptr = + reinterpret_cast( + this->smem_iterator_B_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / + IteratorB::kAccessesPerVector / 8; + + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) { + auto gmem_ptr = iterator_B.get(); + + if (zero_outside_bounds_ || + SharedMemoryClear == SharedMemoryClearOption::kZfill) { + cutlass::arch::cp_async_zfill( + dst_ptr + v, gmem_ptr, iterator_B.valid()); + } else { + cutlass::arch::cp_async( + dst_ptr + v, gmem_ptr, iterator_B.valid()); + } + + ++iterator_B; + } + ++this->smem_iterator_B_; + } + } + } + + template + CUTLASS_DEVICE static void _prologue( + IteratorA& iterator_A, + IteratorB& iterator_B, + int32_t& gemm_k_iterations, + SmemIteratorA& smem_iterator_A_, + SmemIteratorB& smem_iterator_B_) { + // Issue several complete stages + CUTLASS_PRAGMA_UNROLL + for (int stage = 0; stage < kNumStagesConcurrentLoad; + ++stage, --gemm_k_iterations) { + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + + iterator_A.set_iteration_index(0); + smem_iterator_A_.set_iteration_index(0); + + // Async Copy for operand A + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) { + typename IteratorA::AccessType* dst_ptr = + reinterpret_cast( + smem_iterator_A_.get()); + + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { + int const kSrcBytes = + sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / + IteratorA::kAccessesPerVector / 8; + + int src_bytes = (iterator_A.valid() ? kSrcBytes : 0); + + if (kLoadA) { + cutlass::arch::cp_async_zfill( + dst_ptr + v, iterator_A.get(), iterator_A.valid()); + } + + ++iterator_A; + } + + ++smem_iterator_A_; + } + + iterator_B.set_iteration_index(0); + smem_iterator_B_.set_iteration_index(0); + + // Async Copy for operand B + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) { + typename IteratorB::AccessType* dst_ptr = + reinterpret_cast( + smem_iterator_B_.get()); + + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) { + int const kSrcBytes = + sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / + IteratorB::kAccessesPerVector / 8; + + if (kLoadB) { + cutlass::arch::cp_async_zfill( + dst_ptr + v, iterator_B.get(), iterator_B.valid()); + } + + ++iterator_B; + } + + ++smem_iterator_B_; + } + + // Move to the next stage + iterator_A.add_tile_offset({0, 1}); + iterator_B.add_tile_offset({1, 0}); + + smem_iterator_A_.add_tile_offset({0, 1}); + smem_iterator_B_.add_tile_offset({1, 0}); + + // Defines the boundary of a stage of cp.async. + cutlass::arch::cp_async_fence(); + } + } + + /// Perform a threadblock-scoped matrix multiply-accumulate + CUTLASS_DEVICE + void operator()( + ///< problem size of GEMM + int gemm_k_iterations, + ///< destination accumulator tile + FragmentC& accum, + ///< iterator over A operand in global memory + IteratorA iterator_A, + ///< iterator over B operand in global memory + IteratorB iterator_B, + ///< initial value of accumulator + FragmentC const& src_accum) { + // + // Prologue + // + + if (!prologue_done_) { + _prologue( + iterator_A, + iterator_B, + gemm_k_iterations, + smem_iterator_A_, + smem_iterator_B_); + } else if (!kSmemContainsEntireMat) { + _prologue( + iterator_A, + iterator_B, + gemm_k_iterations, + smem_iterator_A_, + smem_iterator_B_); + } else { + gemm_k_iterations -= kNumStagesConcurrentLoad; + } + + // Perform accumulation in the 'd' output operand + accum = src_accum; + + // + // Clear the remaining tiles of SMEM. This is a functional requirement for + // some kernels so that all accumulator elements outside the GEMM footprint + // are zero. + // + + if (SharedMemoryClear == SharedMemoryClearOption::kClearLastStage) { + /// Iterator to write threadblock-scoped tile of A operand to shared + /// memory + SmemIteratorA last_smem_iterator_A(this->smem_iterator_A_); + + typename IteratorA::AccessType zero_A; + zero_A.clear(); + + last_smem_iterator_A.set_iteration_index(0); + + // Async Copy for operand A + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) { + typename IteratorA::AccessType* dst_ptr = + reinterpret_cast( + last_smem_iterator_A.get()); + + *dst_ptr = zero_A; + + ++last_smem_iterator_A; + } + + /// Iterator to write threadblock-scoped tile of B operand to shared + /// memory + SmemIteratorB last_smem_iterator_B(this->smem_iterator_B_); + typename IteratorB::AccessType zero_B; + + zero_B.clear(); + last_smem_iterator_B.set_iteration_index(0); + + // Async Copy for operand B + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) { + typename IteratorB::AccessType* dst_ptr = + reinterpret_cast( + last_smem_iterator_B.get()); + + *dst_ptr = zero_B; + + ++last_smem_iterator_B; + } + } + + // Waits until kStages-2 stages have committed. + cutlass::arch::cp_async_wait(); + __syncthreads(); + + // Pair of fragments used to overlap shared memory loads and math + // instructions + WarpLoadedFragmentA warp_loaded_frag_A[2]; + WarpLoadedFragmentB warp_loaded_frag_B[2]; + WarpTransformedFragmentA warp_transformed_frag_A[2]; + WarpTransformedFragmentB warp_transformed_frag_B[2]; + + Operator warp_mma; + + this->warp_tile_iterator_A_.set_kgroup_index(0); + this->warp_tile_iterator_B_.set_kgroup_index(0); + + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[0]); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B[0]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + + int smem_write_stage_idx = Base::kStages - 1; + int smem_read_stage_idx = 0; + + warp_mma.transform( + warp_transformed_frag_A[0], + warp_transformed_frag_B[0], + warp_loaded_frag_A[0], + warp_loaded_frag_B[0]); + + // tf32x3 kernels use staging accumulation. warp_mma uses a temporary + // accumulator and this temporary accumulator is added to the final + // accumulator once in every mainloop iteration. + plus plus_accum; + + FragmentC tmp_accum; + + if (platform::is_same< + typename Operator::MathOperator, + arch::OpMultiplyAddFastF32>::value || + platform::is_same< + typename Operator::MathOperator, + arch::OpMultiplyAddComplexFastF32>::value) { + tmp_accum.clear(); + } + + // + // Mainloop + // + + CUTLASS_GEMM_LOOP + for (; gemm_k_iterations > (-kNumStagesConcurrentLoad);) { + // + // Loop over GEMM K dimension + // + + // Computes a warp-level GEMM on data held in shared memory + // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate + CUTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; + ++warp_mma_k) { + // Load warp-level tiles from shared memory, wrapping to k offset if + // this is the last group as the case may be. + + this->warp_tile_iterator_A_.set_kgroup_index( + (warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_B_.set_kgroup_index( + (warp_mma_k + 1) % Base::kWarpGemmIterations); + + // In case of a non-circular buffer ("kSmemContainsEntireMat") + // make sure we don't load out of bounds data. + if (!kSmemContainsEntireMat || + gemm_k_iterations > (-kNumStagesConcurrentLoad) || + warp_mma_k < Base::kWarpGemmIterations - 1) { + this->warp_tile_iterator_A_.load( + warp_loaded_frag_A[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_B_.load( + warp_loaded_frag_B[(warp_mma_k + 1) % 2]); + } + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + if (warp_mma_k > 0) + warp_mma.transform( + warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + warp_loaded_frag_A[warp_mma_k % 2], + warp_loaded_frag_B[warp_mma_k % 2]); + + if (platform::is_same< + typename Operator::MathOperator, + arch::OpMultiplyAddFastF32>::value || + platform::is_same< + typename Operator::MathOperator, + arch::OpMultiplyAddComplexFastF32>::value) { + warp_mma( + tmp_accum, + warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + tmp_accum); + + if (warp_mma_k == 0) { + accum = plus_accum(accum, tmp_accum); + tmp_accum.clear(); + } + } else { + warp_mma( + accum, + warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + accum); + } + + // Issue global->shared copies for the this stage + if (!kSmemContainsEntireMat && + warp_mma_k < Base::kWarpGemmIterations - 1) { + int group_start_iteration_A, group_start_iteration_B; + + group_start_iteration_A = warp_mma_k * Detail::kAccessesPerGroupA; + group_start_iteration_B = warp_mma_k * Detail::kAccessesPerGroupB; + + copy_tiles_and_advance( + iterator_A, + iterator_B, + group_start_iteration_A, + group_start_iteration_B); + } + + if (warp_mma_k + 2 == Base::kWarpGemmIterations) { + if (!kSmemContainsEntireMat) { + int group_start_iteration_A, group_start_iteration_B; + group_start_iteration_A = + (warp_mma_k + 1) * Detail::kAccessesPerGroupA; + group_start_iteration_B = + (warp_mma_k + 1) * Detail::kAccessesPerGroupB; + + copy_tiles_and_advance( + iterator_A, + iterator_B, + group_start_iteration_A, + group_start_iteration_B); + } + + // Inserts a memory fence between stages of cp.async instructions. + cutlass::arch::cp_async_fence(); + + // Waits until kStages-2 stages have committed. + cutlass::arch::cp_async_wait(); + __syncthreads(); + + // Move to the next stage + iterator_A.add_tile_offset({0, 1}); + iterator_B.add_tile_offset({1, 0}); + + this->smem_iterator_A_.add_tile_offset({0, 1}); + this->smem_iterator_B_.add_tile_offset({1, 0}); + + // Add negative offsets to return iterators to the 'start' of the + // circular buffer in shared memory + if (smem_write_stage_idx == (Base::kStages - 1)) { + this->smem_iterator_A_.add_tile_offset({0, -Base::kStages}); + this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); + smem_write_stage_idx = 0; + } else { + ++smem_write_stage_idx; + } + + if (!kSmemContainsEntireMat && + smem_read_stage_idx == (Base::kStages - 1)) { + this->warp_tile_iterator_A_.add_tile_offset( + {0, + -Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations}); + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations, + 0}); + smem_read_stage_idx = 0; + } else { + ++smem_read_stage_idx; + } + + --gemm_k_iterations; + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + } + + // Do any conversions feeding the first stage at the end of the loop so + // we can start right away on mma instructions + if (warp_mma_k + 1 == Base::kWarpGemmIterations) + warp_mma.transform( + warp_transformed_frag_A[(warp_mma_k + 1) % 2], + warp_transformed_frag_B[(warp_mma_k + 1) % 2], + warp_loaded_frag_A[(warp_mma_k + 1) % 2], + warp_loaded_frag_B[(warp_mma_k + 1) % 2]); + } + } + + if (platform::is_same< + typename Operator::MathOperator, + arch::OpMultiplyAddFastF32>::value || + platform::is_same< + typename Operator::MathOperator, + arch::OpMultiplyAddComplexFastF32>::value) { + accum = plus_accum(accum, tmp_accum); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fmha/gemm/custom_mma_pipelined.h b/natten/csrc/include/natten/cuda/fmha/gemm/custom_mma_pipelined.h new file mode 100644 index 00000000..af987377 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha/gemm/custom_mma_pipelined.h @@ -0,0 +1,408 @@ +/* + * Copied from xFormers (https://github.com/facebookresearch/xformers/) + * + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + */ +/*************************************************************************************************** + * Copyright (c) 2017 - 2026 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 Template for a double-buffered threadblock-scoped GEMM kernel. +*/ + +#pragma once + +#include +#include +#include +#include + +#include +#include + +#include +#include + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace gemm { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Transformation applied to A operand + typename TransformA_ = NumericArrayConverter< + typename SmemIteratorA_::Element, + typename IteratorA_::Element, + IteratorA_::Fragment::kElements>, + /// + /// Transformation applied to B operand + typename TransformB_ = NumericArrayConverter< + typename SmemIteratorB_::Element, + typename IteratorB_::Element, + IteratorB_::Fragment::kElements>, + /// Used for partial specialization + typename Enable = bool> +class CustomMmaPipelined : public CustomMmaBase { + public: + ///< Base class + using Base = CustomMmaBase; + + using Shape = + Shape_; ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using IteratorA = + IteratorA_; ///< Iterates over tiles of A operand in global memory + using IteratorB = + IteratorB_; ///< Iterates over tiles of B operand in global memory + using ElementC = ElementC_; ///< Data type of accumulator matrix + using LayoutC = LayoutC_; ///< Layout of accumulator matrix + using Policy = Policy_; ///< Policy describing tuning details + + using SmemIteratorA = SmemIteratorA_; + using SmemIteratorB = SmemIteratorB_; + + using TransformA = TransformA_; + using TransformB = TransformB_; + + // + // Dependent types + // + + /// Fragment of operand A loaded from global memory + using FragmentA = typename IteratorA::Fragment; + + /// Fragment of operand B loaded from global memory + using FragmentB = typename IteratorB::Fragment; + + /// Fragment of accumulator tile + using FragmentC = typename Policy::Operator::FragmentC; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Obtain the arch tag from the warp-level operator + using ArchTag = typename Policy::Operator::ArchTag; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = Operator::kTransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = Operator::kTransformB; + + // staticaly assert kStages for MmaPipelined is two (Double-buffered pipeline) + static_assert( + (Base::kStages == 2), + "MmaPipelined requires kStages set to value 2"); + + static bool const kSmemContainsEntireMat = false; + + private: + using WarpFragmentA = typename Operator::FragmentA; + using WarpFragmentB = typename Operator::FragmentB; + + protected: + /// Iterator to write threadblock-scoped tile of A operand to shared memory + SmemIteratorA smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB smem_iterator_B_; + + public: + /// Construct from tensor references + CUTLASS_DEVICE + CustomMmaPipelined( + typename Base::SharedStorageA& shared_storageA, + typename Base::SharedStorageB& shared_storageB, + int thread_idx, ///< ID within the threadblock + int warp_idx, ///< ID of warp + int lane_idx ///< ID of each thread within a warp + ) + : Base(shared_storageA, shared_storageB, thread_idx, warp_idx, lane_idx), + smem_iterator_A_(shared_storageA.ref(), thread_idx), + smem_iterator_B_(shared_storageB.ref(), thread_idx) { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + + int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; + int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset( + {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); + } + CUTLASS_DEVICE + CustomMmaPipelined( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + typename Base::SharedStorage& st, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx) + : CustomMmaPipelined( + st.operand_A, + st.operand_B, + thread_idx, + warp_idx, + lane_idx) {} + + CUTLASS_DEVICE + void set_prologue_done(bool value) { + // NOT IMPLEMENTED FOR PIPELINED + } + + CUTLASS_DEVICE + void set_zero_outside_bounds(bool value) { + // NOT NEEDED FOR PIPELINED + // shared memory will always be zero-filled + } + + template + CUTLASS_DEVICE static void prologue( + typename Base::SharedStorage& shared_storage, + ///< iterator over A operand in global memory + IteratorA iterator_A, + ///< iterator over B operand in global memory + IteratorB iterator_B, + int thread_idx, + int problem_size_k) { + prologue( + shared_storage.operand_A, + shared_storage.operand_B, + iterator_A, + iterator_B, + thread_idx, + problem_size_k); + } + + template + CUTLASS_DEVICE static void prologue( + typename Base::SharedStorageA& shared_storageA, + typename Base::SharedStorageB& shared_storageB, + ///< iterator over A operand in global memory + IteratorA iterator_A, + ///< iterator over B operand in global memory + IteratorB iterator_B, + int thread_idx, + int problem_size_k) { + // NOT IMPLEMENTED FOR PIPELINED + } + + /// Perform a threadblock-scoped matrix multiply-accumulate + CUTLASS_DEVICE + void operator()( + int gemm_k_iterations, ///< number of iterations of the mainloop + FragmentC& accum, ///< destination accumulator tile + IteratorA iterator_A, ///< iterator over A operand in global memory + IteratorB iterator_B, ///< iterator over B operand in global memory + FragmentC const& src_accum, ///< source accumulator tile + TransformA transform_A = + TransformA(), ///< transformation applied to A fragment + TransformB transform_B = + TransformB()) { ///< transformation applied to B fragment + + // + // Prologue + // + + // Perform accumulation in the 'd' output operand + accum = src_accum; + + FragmentA tb_frag_A; + FragmentB tb_frag_B; + + tb_frag_A.clear(); + tb_frag_B.clear(); + + // The last kblock is loaded in the prolog + iterator_A.load(tb_frag_A); + iterator_B.load(tb_frag_B); + + ++iterator_A; + ++iterator_B; + + this->smem_iterator_A_.store(transform_A(tb_frag_A)); + this->smem_iterator_B_.store(transform_B(tb_frag_B)); + + ++this->smem_iterator_A_; + ++this->smem_iterator_B_; + + __syncthreads(); + + // Pair of fragments used to overlap shared memory loads and math + // instructions + WarpFragmentA warp_frag_A[2]; + WarpFragmentB warp_frag_B[2]; + + this->warp_tile_iterator_A_.set_kgroup_index(0); + this->warp_tile_iterator_B_.set_kgroup_index(0); + + this->warp_tile_iterator_A_.load(warp_frag_A[0]); + this->warp_tile_iterator_B_.load(warp_frag_B[0]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + Operator warp_mma; + + int smem_write_stage_idx = 1; + + // Avoid reading out of bounds + iterator_A.clear_mask(gemm_k_iterations <= 1); + iterator_B.clear_mask(gemm_k_iterations <= 1); + + // Issue loads during the first warp-level matrix multiply-add *AFTER* + // issuing shared memory loads (which have the tightest latency + // requirement). + + // + // Mainloop + // + + // Note: The main loop does not support Base::kWarpGemmIterations == 2. + CUTLASS_GEMM_LOOP + for (; gemm_k_iterations > 0; --gemm_k_iterations) { + // + // Loop over GEMM K dimension + // + + CUTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; + ++warp_mma_k) { + // Load warp-level tiles from shared memory, wrapping to k offset if + // this is the last group as the case may be. + + if (warp_mma_k == Base::kWarpGemmIterations - 1) { + // Write fragments to shared memory + this->smem_iterator_A_.store(transform_A(tb_frag_A)); + + this->smem_iterator_B_.store(transform_B(tb_frag_B)); + + __syncthreads(); + + ++this->smem_iterator_A_; + ++this->smem_iterator_B_; + + // Add negative offsets to return iterators to the 'start' of the + // circular buffer in shared memory + if (smem_write_stage_idx == 1) { + this->smem_iterator_A_.add_tile_offset({0, -Base::kStages}); + this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); + } else { + this->warp_tile_iterator_A_.add_tile_offset( + {0, + -Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations}); + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations, + 0}); + } + + smem_write_stage_idx ^= 1; + } + + this->warp_tile_iterator_A_.set_kgroup_index( + (warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_B_.set_kgroup_index( + (warp_mma_k + 1) % Base::kWarpGemmIterations); + + this->warp_tile_iterator_A_.load(warp_frag_A[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_B_.load(warp_frag_B[(warp_mma_k + 1) % 2]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + if (warp_mma_k == 0) { + iterator_A.load(tb_frag_A); + iterator_B.load(tb_frag_B); + + ++iterator_A; + ++iterator_B; + + // Avoid reading out of bounds if this was the last loop iteration + iterator_A.clear_mask(gemm_k_iterations <= 2); + iterator_B.clear_mask(gemm_k_iterations <= 2); + } + + warp_mma( + accum, + warp_frag_A[warp_mma_k % 2], + warp_frag_B[warp_mma_k % 2], + accum); + } + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fmha/gemm/find_default_mma.h b/natten/csrc/include/natten/cuda/fmha/gemm/find_default_mma.h new file mode 100644 index 00000000..f1adad21 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha/gemm/find_default_mma.h @@ -0,0 +1,167 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ +/*! \file + \brief Cutlass provides helper template functions to figure out the right + datastructures to instanciate to run a GEMM with various parameters (see + `cutlass/gemm/threadblock/default_mma.h`). However, due to template + instantiation priority rules, it will only create an MmaMultiStage with + kStages=3 (otherwise creates an MmePipelined - which is not compatible with + FastF32). kStages=3 uses too much shared memory and we want to use kStages=2, + so we just copy-pasted some code from `default_mma.h` and + `default_mma_core.h` files and wrapped this template to allow our usecase. + + This is really only for the FastF32 case - aka using TensorCores with fp32. +*/ + +#pragma once + +#include +#include +#include +#include +#include + +namespace cutlass { +namespace gemm { +namespace threadblock { + +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Layout type for C and D matrix operand + typename LayoutC, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation perfomed by GEMM + typename Operator, + typename Enable_ = void> +struct FindDefaultMma { + static constexpr bool AccumulatorsInRowMajor = false; + static constexpr SharedMemoryClearOption SharedMemoryClear = + SharedMemoryClearOption::kNone; + using DefaultMma = cutlass::gemm::threadblock::DefaultMma< + ElementA, + LayoutA, + kAlignmentA, + ElementB, + LayoutB, + kAlignmentB, + ElementAccumulator, + LayoutC, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + Stages, + Operator, + AccumulatorsInRowMajor, + SharedMemoryClear>; +}; + +/// Specialization for sm80 / FastF32 / multistage with kStages=2 +template < + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + int kStages, + typename Operator> +struct FindDefaultMma< + ElementA_, + LayoutA_, + kAlignmentA, + ElementB_, + LayoutB_, + kAlignmentB, + ElementAccumulator, + layout::RowMajor, + arch::OpClassTensorOp, + arch::Sm80, + ThreadblockShape, + WarpShape, + InstructionShape, + kStages, + Operator, + typename cutlass::platform::enable_if<(kAlignmentA > 1)>::type> { + using LayoutC = layout::RowMajor; + using OperatorClass = arch::OpClassTensorOp; + using ArchTag = arch::Sm80; + + using DefaultMma_ = cutlass::gemm::threadblock::DefaultMma< + ElementA_, + LayoutA_, + kAlignmentA, + ElementB_, + LayoutB_, + kAlignmentB, + ElementAccumulator, + LayoutC, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + 3, + Operator>; + struct DefaultMma : DefaultMma_ { + using MmaCore_ = typename DefaultMma_::MmaCore; + // Define the threadblock-scoped multistage matrix multiply + using ThreadblockMma = cutlass::gemm::threadblock::MmaMultistage< + typename MmaCore_::Shape, + typename DefaultMma_::IteratorA, + typename MmaCore_::SmemIteratorA, + MmaCore_::kCacheOpA, + typename DefaultMma_::IteratorB, + typename MmaCore_::SmemIteratorB, + MmaCore_::kCacheOpB, + ElementAccumulator, + LayoutC, + typename MmaCore_::MmaPolicy, + kStages>; + }; +}; + +} // namespace threadblock +} // namespace gemm +} // namespace cutlass diff --git a/natten/csrc/include/natten/cuda/fmha/gemm/mma_accum_lambda_iterator.h b/natten/csrc/include/natten/cuda/fmha/gemm/mma_accum_lambda_iterator.h new file mode 100644 index 00000000..afc3c9b7 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha/gemm/mma_accum_lambda_iterator.h @@ -0,0 +1,354 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ +#pragma once + +#include +#include +#include +#include +#include + +/* +TensorCores have different accumulator layouts. +This file provides a class to easily map the accumulator +i-th element with the corresponding matrix row/col. +*/ + +template +struct AccumLambdaIteratorSm80 { + static_assert( + cutlass::platform:: + is_same::value, + "only RowMajor is supported"); + + using Policy = typename T::Policy; + using InstructionShape = typename T::InstructionShape; + using OpDelta = typename T::OpDelta; + using Shape = typename T::Shape; + static int const kElementsPerAccess = InstructionShape::kN / 4; + static int const kRowsPerTile = 8; + static int const kAccumulatorRows = InstructionShape::kM / kRowsPerTile; + + static cutlass::MatrixCoord CUTLASS_DEVICE get_lane_offset( + int8_t lane_id, + int8_t warp_id, + typename T::TensorCoord const& tile_offset) { + int quad = (lane_id >> 2); + int lane_in_quad = (lane_id & 3); + return cutlass::MatrixCoord( + quad + tile_offset.row() * Shape::kRow, + lane_in_quad * kElementsPerAccess + + tile_offset.column() * Shape::kColumn); + } + + template + CUTLASS_DEVICE static void iterateRows( + cutlass::MatrixCoord& lane_offset, + FA beginRow, + FB op, + FC endRow) { + // See cutlass/gemm/warp/mma_tensor_op_tile_iterator.h + CUTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + CUTLASS_PRAGMA_UNROLL + for (int row = 0; row < kAccumulatorRows; ++row) { + int accum_m = mma_m * InstructionShape::kM * OpDelta::kRow + + row * kRowsPerTile + lane_offset.row(); + beginRow(accum_m); + + CUTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + int mma_accum_start = kAccumulatorRows * kElementsPerAccess * + (mma_n * Policy::MmaIterations::kRow + mma_m); + CUTLASS_PRAGMA_UNROLL + for (int col = 0; col < kElementsPerAccess; ++col) { + int accum_n = mma_n * InstructionShape::kN * OpDelta::kColumn + + col + lane_offset.column(); + int idx = mma_accum_start + row * kElementsPerAccess + col; + op(accum_m, accum_n, idx); + } + } + + endRow(accum_m); + } + } + } + + template + CUTLASS_DEVICE static bool reduceSameRow(int lane_id, DT& myValue, F fn) { + // In each warp, 4 threads will work on the same row + // - the ones with the same `quad` + auto otherV = __shfl_xor_sync(0xffffffff, myValue, 1); + myValue = fn(myValue, otherV); + otherV = __shfl_xor_sync(0xffffffff, myValue, 2); + myValue = fn(myValue, otherV); + int lane_in_quad = (lane_id & 3); + return lane_in_quad == 0; + } +}; + +template +struct AccumLambdaIteratorSm70 { + static_assert( + cutlass::platform:: + is_same::value, + "only RowMajor is supported"); + + using Policy = typename T::Policy; + using InstructionShape = typename T::InstructionShape; + using OpDelta = typename T::OpDelta; + using Shape = typename T::Shape; + using Element = accum_t; + + static int const kElementsPerPartial = 4; + using EleShapePerPatial = typename cutlass::platform::conditional< + cutlass::platform::is_same::value, + cutlass::MatrixShape<2, 2>, + cutlass::MatrixShape<1, 4>>::type; + static int const kElementsPerMma = 8; + static int const kAccumulatorPatials = 2; + using QuadShapePerPatialMma = cutlass::MatrixShape<4, 4>; + + static cutlass::MatrixCoord CUTLASS_DEVICE get_lane_offset( + int8_t lane_id, + int8_t warp_id, + typename T::TensorCoord const& tile_offset) { + int quad = (lane_id >> 2); + int lane_in_quad = (lane_id & 3); + int accum_m, accum_n; + + if (cutlass::platform::is_same::value) { + // (quad[2],quad[0])+lane_in_quad[0] + accum_m = (((quad & 0x4) >> 1) + (quad & 0x1)) * 8 + (lane_in_quad & 1); + // (quad[1])+lane_in_quad[1] + accum_n = + ((quad >> 1) & 0x1) * kElementsPerPartial * kAccumulatorPatials + + (lane_in_quad & 2); + } else { + accum_m = (((quad & 0x4) >> 1) + (quad & 0x1)) * 8 + + lane_in_quad; // (quad[2],quad[0]) + accum_n = ((quad >> 1) & 0x1) * kElementsPerPartial * kAccumulatorPatials; + } + return cutlass::MatrixCoord( + accum_m + tile_offset.row() * Shape::kRow, + accum_n + tile_offset.column() * Shape::kColumn); + } + + template + CUTLASS_DEVICE static bool reduceSameRow(int lane_id, DT& myValue, F fn) { + static_assert( + cutlass::platform::is_same::value, + "update to support non-float accum"); + // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-fragment-mma-884-f16 + // T0 & T2 share same line within a quad + auto otherV = __shfl_xor_sync(0xffffffff, myValue, 1 << 1); + myValue = fn(myValue, otherV); + // quad 0 and quad 2 are on the same lines + otherV = __shfl_xor_sync(0xffffffff, myValue, 1 << 3); + myValue = fn(myValue, otherV); + return (lane_id & ((1 << 1) | (1 << 3))) == 0; + } + + template + CUTLASS_DEVICE static void iterateRows( + cutlass::MatrixCoord& lane_offset, + FA beginRow, + FB op, + FC endRow) { + CUTLASS_PRAGMA_UNROLL + for (int tile_m = 0; tile_m < Policy::TileIterations::kRow; ++tile_m) { + CUTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + CUTLASS_PRAGMA_UNROLL + for (int m = 0; m < EleShapePerPatial::kRow; ++m) { + int accum_m = tile_m * Policy::InterleavedTile::kRow + + mma_m * QuadShapePerPatialMma::kRow + m * 2 + lane_offset.row(); + beginRow(accum_m); + + CUTLASS_PRAGMA_UNROLL + for (int tile_n = 0; tile_n < Policy::TileIterations::kColumn; + ++tile_n) { + CUTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; + ++mma_n) { + CUTLASS_PRAGMA_UNROLL + for (int p = 0; p < kAccumulatorPatials; ++p) { + CUTLASS_PRAGMA_UNROLL + for (int n = 0; n < EleShapePerPatial::kColumn; ++n) { + int mma_accum_start = + (((tile_n * Policy::TileIterations::kRow + tile_m) * + Policy::MmaIterations::kColumn + + mma_n) * + Policy::MmaIterations::kRow + + mma_m) * + kElementsPerMma; + int accum_n = tile_n * Policy::InterleavedTile::kColumn + + mma_n * QuadShapePerPatialMma::kColumn + + p * Policy::InterleavedTile::kColumn / 2 + n + + lane_offset.column(); + int idx = mma_accum_start + p * kElementsPerPartial + + m * EleShapePerPatial::kColumn + n; + op(accum_m, accum_n, idx); + } + } + } + } + endRow(accum_m); + } + } + } + } +}; + +template +struct AccumLambdaIteratorSimt { + using Policy = typename T::Policy; + using Iterations = typename T::Iterations; + using Element = typename T::Element; + using Delta = typename T::Delta; + using Shape = typename T::Shape; + static_assert( + cutlass::platform:: + is_same::value, + "only RowMajor is supported"); + + template + CUTLASS_DEVICE static bool reduceSameRow(int lane_id, DT& myValue, F fn) { + CUTLASS_PRAGMA_UNROLL + for (int bit = 1; bit < Policy::WarpShape::kColumn; bit *= 2) { + auto otherV = __shfl_xor_sync(0xffffffff, myValue, bit); + myValue = fn(myValue, otherV); + } + return (lane_id & (Policy::WarpShape::kColumn - 1)) == 0; + } + + template + CUTLASS_DEVICE static void iterateRows( + cutlass::MatrixCoord& lane_offset, + FA beginRow, + FB op, + FC endRow) { + CUTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Iterations::kRow; ++mma_m) { + CUTLASS_PRAGMA_UNROLL + for (int m = 0; m < Policy::LaneMmaShape::kM; ++m) { + int accum_m = mma_m * Delta::kRow + m + lane_offset.row(); + beginRow(accum_m); + + CUTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Iterations::kColumn; ++mma_n) { + int accum_n = + mma_n * Policy::WarpShape::kColumn * Policy::LaneMmaShape::kN + + lane_offset.column(); + CUTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::LaneMmaShape::kN; ++n) { + int idx = n + + Policy::LaneMmaShape::kN * + (mma_n + + Iterations::kColumn * + (m + mma_m * Policy::LaneMmaShape::kM)); + op(accum_m, accum_n + n, idx); + } + } + endRow(accum_m); + } + } + } + + static cutlass::MatrixCoord CUTLASS_DEVICE get_lane_offset( + int8_t lane_id, + int8_t warp_id, + typename T::TensorCoord const& tile_offset) { + static_assert( + cutlass::platform::is_same< + typename Policy::LaneLayout, + cutlass::layout::RowMajorInterleaved<1>>::value, + ""); + typename Policy::LaneLayout lane_layout = Policy::get_lane_layout(); + + cutlass::MatrixCoord lane_offset = lane_layout.inverse(lane_id) * + cutlass::MatrixCoord(Policy::LaneMmaShape::kM, + Policy::LaneMmaShape::kN); + return lane_offset + + tile_offset * cutlass::MatrixCoord(Shape::kRow, Shape::kColumn); + } +}; + +template +struct DefaultMmaAccumLambdaIterator; + +// Simt +template +struct DefaultMmaAccumLambdaIterator< + cutlass::gemm::warp::MmaSimtTileIterator< + S, + cutlass::gemm::Operand::kC, + accum_t, + cutlass::layout::RowMajor, + P, + 1, + 1>, + accum_t, + kWarpSize> { + using WarpIterator = typename cutlass::gemm::warp::MmaSimtTileIterator< + S, + cutlass::gemm::Operand::kC, + accum_t, + cutlass::layout::RowMajor, + P, + 1, + 1>; + using Iterator = AccumLambdaIteratorSimt; +}; + +// TensorOp - Volta +template +struct DefaultMmaAccumLambdaIterator< + cutlass::gemm::warp::MmaVoltaTensorOpAccumulatorTileIterator< + S1, + accum_t, + cutlass::layout::RowMajor, + S2, + cutlass::MatrixShape<1, 1>>, + accum_t, + kWarpSize> { + using WarpIterator = + typename cutlass::gemm::warp::MmaVoltaTensorOpAccumulatorTileIterator< + S1, + accum_t, + cutlass::layout::RowMajor, + S2, + cutlass::MatrixShape<1, 1>>; + using Iterator = AccumLambdaIteratorSm70; +}; + +// TensorOp - Sm75+ +template < + typename S1, + typename S2, + typename S3, + typename accum_t, + int kWarpSize> +struct DefaultMmaAccumLambdaIterator< + cutlass::gemm::warp::MmaTensorOpAccumulatorTileIterator< + S1, + accum_t, + cutlass::layout::RowMajor, + S2, + S3>, + accum_t, + kWarpSize> { + using WarpIterator = + typename cutlass::gemm::warp::MmaTensorOpAccumulatorTileIterator< + S1, + accum_t, + cutlass::layout::RowMajor, + S2, + S3>; + using Iterator = AccumLambdaIteratorSm80; +}; diff --git a/natten/csrc/include/natten/cuda/fmha/gemm/mma_from_smem.h b/natten/csrc/include/natten/cuda/fmha/gemm/mma_from_smem.h new file mode 100644 index 00000000..23cb4621 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha/gemm/mma_from_smem.h @@ -0,0 +1,1916 @@ +/* + * Copied from xFormers (https://github.com/facebookresearch/xformers/) and + * edited. + * + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + */ +/*************************************************************************************************** + * Copyright (c) 2017 - 2026 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 Tools and utils to store a GEMM output in shmem, and to use that + output as operandA for another GEMM back-to-back +*/ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace cutlass { +namespace gemm { +namespace threadblock { + +/// Shared storage object needed by accumulator +/// From 13_two_tensor_op_fusion/threadblock/b2b_mma_base_smem_accumulator.h +template < + typename Shape_, + typename Element_, + typename Layout_, + typename Padding_> +class AccumulatorSharedStorage { + public: + // + // Type definitions + // + using Shape = Shape_; + using Element = Element_; + using Layout = Layout_; + using Padding = Padding_; + + /// Tensor reference to the accumulator + using TensorRefAccum = cutlass::TensorRef; + + /// Shape of the accumulator matrix in shared memory + using ShapeAccum = cutlass:: + MatrixShape; + + public: + // + // Data members + // + + /// Buffer for accumulator + cutlass::AlignedBuffer accum; + + public: + // + // Methods + // + + /// Returns a layout object for the Accum matrix + CUTLASS_DEVICE + static Layout LayoutAccum() { + return Layout::packed({ShapeAccum::kRow, ShapeAccum::kColumn}); + } + + /// Returns a TensorRef to the Accumulator + CUTLASS_HOST_DEVICE + TensorRefAccum accum_ref() { + return TensorRefAccum{accum.data(), LayoutAccum()}; + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// Taken from +// https://github.com/NVIDIA/cutlass/blob/master/examples/13_two_tensor_op_fusion/threadblock/b2b_mma_base_smem_accumulator.h +//////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + // Maximum K dimension - also the dimension of the shared-memory + // holding `OperandA` + int kMaxK_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Layout in shared-memory of operand A + typename SmemLayoutA, + /// Used for partial specialization + typename Enable = bool> +class MmaBaseFromSharedMemory { + public: + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + static constexpr int kMaxK = kMaxK_; + + ///< Policy describing tuning details + using Policy = Policy_; + + // + // Dependent types + // + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Shape describing the overall GEMM computed from shared memory + /// by each warp. + using WarpGemm = typename Policy::Operator::Shape; + + /// Shape describing the number of warps filling the CTA + using WarpCount = GemmShape< + Shape::kM / WarpGemm::kM, + Shape::kN / WarpGemm::kN, + Shape::kK / WarpGemm::kK>; + using WarpCount1 = WarpCount; + + /// Number of warp-level GEMM oeprations + static int const kWarpGemmIterations = + (WarpGemm::kK / Operator::Policy::MmaShape::kK); + static int const kWarpGemmIterations1 = kWarpGemmIterations; + + /// Number of stages + static int const kStages = Stages; + + /// If this is true, we fill the entire shmem buffer at start + /// and don't need to iterate through it in a circular fashion + static bool const kSmemContainsEntireB = kMaxK <= Shape::kK * kStages; + + /// Tensor reference to the A operand + using TensorRefA = TensorRef; + + /// Tensor reference to the B operand + using TensorRefB = + TensorRef; + + // + // Nested structs + // + + /// Shared storage object needed by threadblock-scoped GEMM + class SharedStorage { + public: + // + // Type definitions + // + + /// Shape of the B matrix operand in shared memory + using ShapeB = MatrixShape< + Shape::kK * kStages + Policy::SmemPaddingB::kRow, + Shape::kN + Policy::SmemPaddingB::kColumn>; + + public: + // + // Data members + // + + /// Buffer for B operand + AlignedBuffer operand_B; + + public: + // + // Methods + // + + /// Returns a layout object for the B matrix + CUTLASS_HOST_DEVICE + static typename Operator::LayoutB LayoutB() { + return Operator::LayoutB::packed({ShapeB::kRow, ShapeB::kColumn}); + } + + /// Returns a TensorRef to the B operand + CUTLASS_HOST_DEVICE + TensorRefB operand_B_ref() { + return TensorRefB{operand_B.data(), LayoutB()}; + } + }; + + protected: + // + // Data members + // + + // /// Iterator to load a warp-scoped tile of A operand from shared memory + // typename Operator::IteratorA warp_tile_iterator_A_; + + /// Iterator to load a warp-scoped tile of B operand from shared memory + typename Operator::IteratorB warp_tile_iterator_B_; + + public: + /// Construct from tensor references + CUTLASS_DEVICE + MmaBaseFromSharedMemory( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + TensorRefB& b_tile, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx) + : warp_tile_iterator_B_(b_tile, lane_idx) {} +}; + +namespace { + +// has necessary trait compliance with WarpIteratorFromSmem but doesn't do +// anything, can be default initialized, and uses fragment that takes up +// (almost) no space. this warp iterator is selected at compile time when +// elementwise on-the-fly scaling for operand A is disabled, in which case +// operations related to loading scale factors for operand A get wiped out by +// the compiler. +template +class NoOpWarpIteratorScale { + public: + // in pipelined+multistage MMA implementations we keep an array of fragments. + // if we aren't using scaling we don't want to waste registers on fragments + // of scale elements, so ideally this would be sized 0. + // Since arrays of zero-sized objects are not allowed, using size as 1. + // The compiler will most likely wipe it out anyways. + using Fragment = cutlass::Array; + + CUTLASS_HOST_DEVICE + NoOpWarpIteratorScale() {} + + CUTLASS_HOST_DEVICE + NoOpWarpIteratorScale(TensorRef const&, int) {} + + CUTLASS_HOST_DEVICE + NoOpWarpIteratorScale& add_tile_offset( + typename TensorRef::TensorCoord const&) { + return *this; + } + + CUTLASS_HOST_DEVICE + NoOpWarpIteratorScale& operator++() { + return *this; + } + + CUTLASS_DEVICE + void load(Fragment&) const {} +}; + +// if scaling is enabled, performs fragment elementwise multiplication between +// fragment and its scaling factor. +template +class FragmentElementwiseScaler; + +// specialization for scaling being enabled. +template +class FragmentElementwiseScaler { + public: + // cast scale_frag to correct type then apply elementwise to fragment + CUTLASS_DEVICE + static Fragment apply(Fragment frag, FragmentScale const& scale_frag) { + Fragment converted_scale_frag = cutlass::NumericArrayConverter< + typename Fragment::Element, + typename FragmentScale::Element, + FragmentScale::kElements>()(scale_frag); + return cutlass::multiplies()(frag, converted_scale_frag); + } +}; + +// specialization for scaling being disabled. doesn't do anything and should +// just get wiped out by the compiler. +template +class FragmentElementwiseScaler { + public: + CUTLASS_DEVICE + static Fragment apply(Fragment frag, FragmentScale const&) { + return frag; + } +}; +} // namespace + +//////////////////////////////////////////////////////////////////////////////// +// Taken from +// https://github.com/NVIDIA/cutlass/blob/master/examples/13_two_tensor_op_fusion/threadblock/b2b_mma_pipelined_smem_accumulator.h +//////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + // BEGIN smem + /// Iterates over the intermediate accumulator tile in shared memory + typename WarpIteratorA_, + /// whether or not to perform elementwise multiplication of A + // by another matrix (A_scale) that is also kept in shared memory prior + // to matmul A @ B + bool ScaleOperandA_, + /// Max GEMM problem size in K dimension + int MaxK, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Transformation applied to B operand + typename TransformB_ = NumericArrayConverter< + typename SmemIteratorB_::Element, + typename IteratorB_::Element, + IteratorB_::Fragment::kElements>, + /// Used for partial specialization + typename Enable = bool> +class MmaPipelinedFromSharedMemory : public MmaBaseFromSharedMemory< + Shape_, + MaxK, + Policy_, + 2, + typename WarpIteratorA_::Layout> { + public: + ///< Base class + using Base = MmaBaseFromSharedMemory< + Shape_, + MaxK, + Policy_, + 2, + typename WarpIteratorA_::Layout>; + + using Shape = + Shape_; ///< Size of the Gemm problem - concept: gemm::GemmShape<> + static constexpr bool ScaleOperandA = ScaleOperandA_; + + using WarpIteratorA = WarpIteratorA_; + ///< loads fragments of A_scale from shared memory if operand A scaling is + ///< enabled. otherwise no-op. + using WarpIteratorAScale = typename cutlass::platform::conditional< + ScaleOperandA, + WarpIteratorA, + NoOpWarpIteratorScale>::type; + + using IteratorB = + IteratorB_; ///< Iterates over tiles of B operand in global memory + using ElementC = ElementC_; ///< Data type of accumulator matrix + using LayoutC = LayoutC_; ///< Layout of accumulator matrix + using Policy = Policy_; ///< Policy describing tuning details + + using SmemIteratorB = SmemIteratorB_; + + using TransformB = TransformB_; + + // + // Dependent types + // + + /// Fragment of operand B loaded from global memory + using FragmentB = typename IteratorB::Fragment; + + /// Fragment of accumulator tile + using FragmentC = typename Policy::Operator::FragmentC; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Obtain the arch tag from the warp-level operator + using ArchTag = typename Policy::Operator::ArchTag; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = Operator::kTransformB; + + // staticaly assert kStages for MmaPipelined is two (Double-buffered pipeline) + static_assert( + (Base::kStages == 2), + "MmaPipelined requires kStages set to value 2"); + + private: + using WarpFragmentA = typename Operator::FragmentA; + + /// fragment type of OperandA elementwise scaling matrix. (almost) empty + /// if operand A scaling is disabled. + using WarpFragmentAScale = typename WarpIteratorAScale::Fragment; + + using WarpFragmentB = typename Operator::FragmentB; + + /// applies scaling factor to operand A fragment if operand A scaling is + /// enabled. otherwise no-op. + using FragmentAScaler = FragmentElementwiseScaler< + WarpFragmentA, + WarpFragmentAScale, + ScaleOperandA>; + + protected: + // /// Iterator to write threadblock-scoped tile of A operand to shared memory + // SmemIteratorA smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB smem_iterator_B_; + + /// Iterator to load a warp-scoped tile of A operand from intermediate + /// accumulator tile + WarpIteratorA warp_tile_iterator_A_; + + /// Iterator to load a warp-scoped tile of A_scale from intermediate + /// accumulator tile (only used if ScaleOperandA_ is true) + WarpIteratorAScale warp_tile_iterator_A_scale_; + + public: + /// constructor for MMA with operand A scaling enabled. + CUTLASS_DEVICE + MmaPipelinedFromSharedMemory( + typename Base::TensorRefA a, // Operand A in shared memory + typename Base::TensorRefA a_scale, // Operand A_scale in shared memory + typename Base::TensorRefB + b_staging, // staging memory for loading tiles of B + int thread_idx, + int warp_idx, + int lane_idx) + : Base(b_staging, thread_idx, warp_idx, lane_idx), + warp_tile_iterator_A_(a, lane_idx), + warp_tile_iterator_A_scale_(a_scale, lane_idx), + smem_iterator_B_(b_staging, thread_idx) { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; + int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset( + {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_A_scale_.add_tile_offset( + {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); + } + + /// Construct from tensor references + CUTLASS_DEVICE + MmaPipelinedFromSharedMemory( + typename Base::TensorRefA a, ///< Operand A in shared memory + typename Base::TensorRefB b_staging, ///< staging memory for loading B + int thread_idx, ///< ID within the threadblock + int warp_idx, ///< ID of warp + int lane_idx) ///< ID of each thread within a warp + : Base(b_staging, thread_idx, warp_idx, lane_idx), + warp_tile_iterator_A_(a, lane_idx), + smem_iterator_B_(b_staging, thread_idx) { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + + int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; + int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset( + {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); + } + + // For API compatibility with MmaMultistageFromSharedMemory + // but not supported as it worsens perf: older gpus < sm80 don't + // support async tranfers and have to waste registers + CUTLASS_DEVICE + void set_prologue_done(bool value) {} + CUTLASS_DEVICE + static void prologue( + typename Base::SharedStorage& shared_storage, + IteratorB iterator_B1, + int thread_idx, + int problem_size_0_n) {} + + CUTLASS_DEVICE + static void drain_cp_asyncs() {} + + /// Perform a threadblock-scoped matrix multiply-accumulate + CUTLASS_DEVICE + void operator()( + int gemm_k_iterations, ///< number of iterations of the mainloop + FragmentC& accum, ///< destination accumulator tile + // IteratorA iterator_A, ///< iterator over A + // operand in global memory + IteratorB iterator_B, ///< iterator over B operand in global memory + FragmentC const& src_accum, ///< source accumulator tile + // TransformA transform_A = TransformA(), ///< transformation + // applied to A fragment + TransformB transform_B = + TransformB()) { ///< transformation applied to B fragment + + // + // Prologue + // + + // Perform accumulation in the 'd' output operand + accum = src_accum; + + FragmentB tb_frag_B; + + tb_frag_B.clear(); + + // The last kblock is loaded in the prolog + iterator_B.set_residual_tile(gemm_k_iterations == 1); + iterator_B.load(tb_frag_B); + + ++iterator_B; + + this->smem_iterator_B_.store(transform_B(tb_frag_B)); + + ++this->smem_iterator_B_; + + __syncthreads(); + + // remember that WarpFragmentAScale and WarpIteratorAScale are empty/no-op + // if scaling is disabled. + + // Pair of fragments used to overlap shared memory loads and math + // instructions + WarpFragmentA warp_frag_A[2]; + WarpFragmentAScale warp_frag_A_scale[2]; + WarpFragmentB warp_frag_B[2]; + warp_frag_A[0].clear(); + warp_frag_A_scale[0].clear(); + warp_frag_B[0].clear(); + + this->warp_tile_iterator_B_.set_kgroup_index(0); + + this->warp_tile_iterator_A_.load(warp_frag_A[0]); + this->warp_tile_iterator_A_scale_.load(warp_frag_A_scale[0]); + this->warp_tile_iterator_B_.load(warp_frag_B[0]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_A_scale_; + ++this->warp_tile_iterator_B_; + + Operator warp_mma; + + int smem_write_stage_idx = 1; + + // Avoid reading out of bounds + iterator_B.set_residual_tile(gemm_k_iterations == 2); + iterator_B.clear_mask(gemm_k_iterations <= 1); + + // Issue loads during the first warp-level matrix multiply-add *AFTER* + // issuing shared memory loads (which have the tightest latency + // requirement). + + // + // Mainloop + // + + // Note: The main loop does not support Base::kWarpGemmIterations == 2. + CUTLASS_GEMM_LOOP + for (; gemm_k_iterations > 0; --gemm_k_iterations) { + // + // Loop over GEMM K dimension + // + + CUTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; + ++warp_mma_k) { + // Load warp-level tiles from shared memory, wrapping to k offset if + // this is the last group as the case may be. + bool hasNext = true; + + if (warp_mma_k == Base::kWarpGemmIterations - 1) { + if (gemm_k_iterations > 1) { + // Write fragments to shared memory + this->smem_iterator_B_.store(transform_B(tb_frag_B)); + } + + __syncthreads(); + + ++this->smem_iterator_B_; + + // Add negative offsets to return iterators to the 'start' of the + // circular buffer in shared memory SMEM: Don't reset iterator A, as + // we are continuing our iteration at this point + if (smem_write_stage_idx == 1) { + this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); + } else { + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations, + 0}); + } + + smem_write_stage_idx ^= 1; + hasNext = gemm_k_iterations > 1; + } + + // Only read the next if we need to + if (hasNext) { + this->warp_tile_iterator_B_.set_kgroup_index( + (warp_mma_k + 1) % Base::kWarpGemmIterations); + + this->warp_tile_iterator_A_.load(warp_frag_A[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_A_scale_.load( + warp_frag_A_scale[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_B_.load(warp_frag_B[(warp_mma_k + 1) % 2]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_A_scale_; + ++this->warp_tile_iterator_B_; + + if (warp_mma_k == 0) { + iterator_B.load(tb_frag_B); + + ++iterator_B; + + // Avoid reading out of bounds if this was the last loop iteration + iterator_B.set_residual_tile(gemm_k_iterations == 3); + iterator_B.clear_mask(gemm_k_iterations <= 2); + } + } + + warp_mma( + accum, + FragmentAScaler::apply( + warp_frag_A[warp_mma_k % 2], warp_frag_A_scale[warp_mma_k % 2]), + warp_frag_B[warp_mma_k % 2], + accum); + } + } + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// Taken from +// https://github.com/NVIDIA/cutlass/blob/master/examples/13_two_tensor_op_fusion/threadblock/b2b_mma_multistage_smem_accumulator.h +//////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape1_, + /// Iterates over the intermediate accumulator tile in shared memory + typename WarpIteratorA1_, + /// whether or not to perform elementwise multiplication of A + // by another matrix (A_scale) that is also kept in shared memory prior + // to matmul A @ B + bool ScaleOperandA_, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB1_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB1_, + /// Cache operation for operand B + cutlass::arch::CacheOperation::Kind CacheOpB1, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy1_, + /// Number of stages, + int Stages_, + int kMaxK_, + /// Used for partial specialization + typename Enable = bool> +class MmaMultistageFromSharedMemory : public MmaBaseFromSharedMemory< + Shape1_, + kMaxK_, + Policy1_, + Stages_, + typename WarpIteratorA1_::Layout> { + public: + ///< Base class + using Base = MmaBaseFromSharedMemory< + Shape1_, + kMaxK_, + Policy1_, + Stages_, + typename WarpIteratorA1_::Layout>; + + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape1 = Shape1_; + ///< Iterates over tiles of B operand in global memory + using IteratorB1 = IteratorB1_; + using IteratorB = IteratorB1; + ///< Policy describing tuning details + using Policy1 = Policy1_; + + using SmemIteratorB1 = SmemIteratorB1_; + using WarpIteratorA1 = WarpIteratorA1_; ///< Iterates over the intermediate + ///< accumulator tile in shared memory + static constexpr bool ScaleOperandA = ScaleOperandA_; + + ///< warp level iterator over A_scale matrix tile kept in shared memory. + ///< if elementwise A scaling is disabled then everything this does is no-op. + using WarpIteratorAScale = typename cutlass::platform::conditional< + ScaleOperandA, + WarpIteratorA1, + NoOpWarpIteratorScale>::type; + ///< Data type of accumulator matrix + using ElementC = ElementC_; + ///< Layout of accumulator matrix + using LayoutC = LayoutC_; + + static cutlass::arch::CacheOperation::Kind const kCacheOpB1 = CacheOpB1; + static constexpr bool kSmemContainsEntireB = Base::kSmemContainsEntireB; + + // + // Dependent types + // + + /// Fragment of accumulator tile + using FragmentC1 = typename Policy1::Operator::FragmentC; + using FragmentC = FragmentC1; + + /// Warp-level Mma + using Operator1 = typename Policy1::Operator; + + /// Minimum architecture is Sm80 to support cp.async + using ArchTag = arch::Sm80; + + /// Complex transform on B operand + static ComplexTransform const kTransformB1 = Operator1::kTransformB; + + /// Internal structure exposed for introspection. + struct Detail { + static_assert( + Base::kWarpGemmIterations1 > 1, + "The pipelined structure requires at least two warp-level " + "GEMM operations."); + + /// Number of cp.async instructions to load one stage of operand B + static int const TBLoadIterationsB1 = + IteratorB1::ThreadMap::Iterations::kCount; + + /// Number of cp.async instructions to load on group of operand B + static int const kAccessesPerGroupB1 = + (TBLoadIterationsB1 + Base::kWarpGemmIterations1 - 1) / + Base::kWarpGemmIterations1; + }; + + static constexpr int kNumStagesConcurrentLoad = + kSmemContainsEntireB ? Base::kStages : Base::kStages - 1; + + private: + using WarpLoadedFragmentA1 = typename Operator1::FragmentA; + /// fragment of OperandA scale matrix. if operand A scaling is disabled this + /// is (almost) empty. + using WarpLoadedFragmentA1Scale = typename WarpIteratorAScale::Fragment; + using WarpLoadedFragmentB1 = typename Operator1::FragmentB; + using WarpTransformedFragmentA1 = typename Operator1::TransformedFragmentA; + using WarpTransformedFragmentB1 = typename Operator1::TransformedFragmentB; + + /// applies elementwise scaling to fragment of A. if operand A scaling is + /// disabled this is a no-op. + using FragmentAScaler = FragmentElementwiseScaler< + WarpLoadedFragmentA1, + WarpLoadedFragmentA1Scale, + ScaleOperandA>; + + private: + // + // Data members + // + + /// Iterator to load a warp-scoped tile of A1 operand from intermediate + /// accumulator tile + WarpIteratorA1 warp_tile_iterator_A1_; + + /// Iterator to load a warp-scoped tile of A1_scale operand from shared memory + /// if operand A scaling is disabled everything this does is a no-op. + WarpIteratorAScale warp_tile_iterator_A1_scale_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB1 smem_iterator_B1_; + + bool prologue_done_; + + public: + /// constructor for MMA with operand A scaling enabled. + CUTLASS_DEVICE + MmaMultistageFromSharedMemory( + typename Base::TensorRefA a, + typename Base::TensorRefA a_scale, + typename Base::TensorRefB b_tile, + int thread_idx, + int warp_idx, + int lane_idx) + : Base(b_tile, thread_idx, warp_idx, lane_idx), + warp_tile_iterator_A1_(a, lane_idx), + warp_tile_iterator_A1_scale_(a_scale, lane_idx), + smem_iterator_B1_(b_tile, thread_idx), + prologue_done_(false) { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + int warp_idx_mn_1 = + warp_idx % (Base::WarpCount1::kM * Base::WarpCount1::kN); + int warp_idx_k_1 = warp_idx / (Base::WarpCount1::kM * Base::WarpCount1::kN); + int warp_idx_m_1 = warp_idx_mn_1 % Base::WarpCount1::kM; + int warp_idx_n_1 = warp_idx_mn_1 / Base::WarpCount1::kM; + + // Add per-warp offsets in units of warp-level tiles + warp_tile_iterator_A1_.add_tile_offset( + {warp_idx_m_1, Base::kWarpGemmIterations1 * warp_idx_k_1}); + warp_tile_iterator_A1_scale_.add_tile_offset( + {warp_idx_m_1, Base::kWarpGemmIterations1 * warp_idx_k_1}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations1 * warp_idx_k_1, warp_idx_n_1}); + } + + /// Construct from tensor references + CUTLASS_DEVICE + MmaMultistageFromSharedMemory( + typename Base::TensorRefA a, + typename Base::TensorRefB b_tile, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx) + : Base(b_tile, thread_idx, warp_idx, lane_idx), + warp_tile_iterator_A1_(a, lane_idx), + smem_iterator_B1_(b_tile, thread_idx), + prologue_done_(false) { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn_1 = + warp_idx % (Base::WarpCount1::kM * Base::WarpCount1::kN); + int warp_idx_k_1 = warp_idx / (Base::WarpCount1::kM * Base::WarpCount1::kN); + + int warp_idx_m_1 = warp_idx_mn_1 % Base::WarpCount1::kM; + int warp_idx_n_1 = warp_idx_mn_1 / Base::WarpCount1::kM; + + // Add per-warp offsets in units of warp-level tiles + warp_tile_iterator_A1_.add_tile_offset( + {warp_idx_m_1, Base::kWarpGemmIterations1 * warp_idx_k_1}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations1 * warp_idx_k_1, warp_idx_n_1}); + } + + CUTLASS_DEVICE + void set_prologue_done(bool value) { + prologue_done_ = value; + } + + CUTLASS_DEVICE + static void prologue( + typename Base::SharedStorage& shared_storage, + IteratorB iterator_B1, + int thread_idx, + int problem_size_0_n) { + SmemIteratorB1 smem_iterator_B1(shared_storage.operand_B_ref(), thread_idx); + _prologue( + iterator_B1, + (problem_size_0_n + Base::Shape::kK - 1) / Base::Shape::kK, + smem_iterator_B1); + } + + CUTLASS_DEVICE + static void drain_cp_asyncs() { + // commit and drain all pending and predicated cp.async pnz from the GEMM + // mainloop + cutlass::arch::cp_async_fence(); + cutlass::arch::cp_async_wait<0>(); + __syncthreads(); + } + + CUTLASS_DEVICE + void copy_tiles_and_advance_1( + IteratorB1& iterator_B1, + int group_start_B1 = 0) { + iterator_B1.set_iteration_index( + group_start_B1 * IteratorB1::kAccessesPerVector); + this->smem_iterator_B1_.set_iteration_index(group_start_B1); + + // Load for operand B + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupB1; ++j) { + if (group_start_B1 + j < Detail::TBLoadIterationsB1) { + typename IteratorB1::AccessType* dst_ptr = + reinterpret_cast( + this->smem_iterator_B1_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorB1::ThreadMap::kElementsPerAccess / + IteratorB1::kAccessesPerVector / 8; + + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB1::kAccessesPerVector; ++v) { + auto gmem_ptr = iterator_B1.get(); + + cutlass::arch::cp_async_zfill( + dst_ptr + v, gmem_ptr, iterator_B1.valid()); + + ++iterator_B1; + } + ++this->smem_iterator_B1_; + } + } + } + + CUTLASS_DEVICE + static void _prologue( + IteratorB& iterator_B1, + int32_t gemm_k_iterations_1, + SmemIteratorB1& smem_iterator_B1_) { + // Issue several complete stages + CUTLASS_PRAGMA_UNROLL + for (int stage = 0; stage < kNumStagesConcurrentLoad; + ++stage, --gemm_k_iterations_1) { + iterator_B1.set_residual_tile(gemm_k_iterations_1 == 1); + iterator_B1.clear_mask(gemm_k_iterations_1 == 0); + + iterator_B1.set_iteration_index(0); + smem_iterator_B1_.set_iteration_index(0); + + // Load for operand B + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::TBLoadIterationsB1; ++j) { + typename IteratorB1::AccessType* dst_ptr = + reinterpret_cast( + smem_iterator_B1_.get()); + + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB1::kAccessesPerVector; ++v) { + int const kSrcBytes = + sizeof_bits::value * + IteratorB1::ThreadMap::kElementsPerAccess / + IteratorB1::kAccessesPerVector / 8; + + cutlass::arch::cp_async_zfill( + dst_ptr + v, iterator_B1.get(), iterator_B1.valid()); + + ++iterator_B1; + } + + ++smem_iterator_B1_; + } + + // Move to the next stage + iterator_B1.add_tile_offset({1, 0}); + + smem_iterator_B1_.add_tile_offset({1, 0}); + + // Defines the boundary of a stage of cp.async. + cutlass::arch::cp_async_fence(); + } + iterator_B1.set_residual_tile(gemm_k_iterations_1 == 1); + iterator_B1.clear_mask(gemm_k_iterations_1 == 0); + } + + /// Perform a threadblock-scoped matrix multiply-accumulate + CUTLASS_DEVICE + void operator()( + ///< problem size of GEMM + int gemm_k_iterations_1_, + ///< destination accumulator tile + FragmentC1& accum, + ///< iterator over B1 operand in global memory + IteratorB1 iterator_B1, + ///< initial value of accumulator + FragmentC1 const& src_accum) { + // 2nd Gemm + + // + // Prologue + // + // Perform accumulation in the 'd' output operand + accum = src_accum; + + if (!prologue_done_) { + _prologue(iterator_B1, gemm_k_iterations_1_, smem_iterator_B1_); + } else if (!kSmemContainsEntireB) { + // Restore the iterators increments + + int gemm_k_iterations_1 = gemm_k_iterations_1_; + // Issue several complete stages + CUTLASS_PRAGMA_UNROLL + for (int stage = 0; stage < kNumStagesConcurrentLoad; + ++stage, --gemm_k_iterations_1) { + iterator_B1.set_iteration_index(0); + this->smem_iterator_B1_.set_iteration_index(0); + + // Load for operand B + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::TBLoadIterationsB1; ++j) { + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB1::kAccessesPerVector; ++v) { + ++iterator_B1; + } + ++this->smem_iterator_B1_; + } + iterator_B1.add_tile_offset({1, 0}); + this->smem_iterator_B1_.add_tile_offset({1, 0}); + } + iterator_B1.set_residual_tile(gemm_k_iterations_1 <= 1); + iterator_B1.clear_mask(gemm_k_iterations_1 <= 0); + } + + // DEPBAR+SYNC + cutlass::arch::cp_async_wait(); + __syncthreads(); + + // remember that WarpFragmentAScale and WarpIteratorAScale are no-op/empty + // if scaling is disabled. + + // Pair of fragments used to overlap shared memory loads and math + // instructions + WarpLoadedFragmentA1 warp_loaded_frag_A1[2]; + WarpLoadedFragmentA1Scale warp_loaded_frag_A1_scale[2]; + WarpLoadedFragmentB1 warp_loaded_frag_B1[2]; + WarpTransformedFragmentA1 warp_transformed_frag_A1[2]; + WarpTransformedFragmentB1 warp_transformed_frag_B1[2]; + + Operator1 warp_mma1; + + warp_tile_iterator_A1_.load(warp_loaded_frag_A1[0]); + ++warp_tile_iterator_A1_; + + warp_tile_iterator_A1_scale_.load(warp_loaded_frag_A1_scale[0]); + ++warp_tile_iterator_A1_scale_; + + this->warp_tile_iterator_B_.set_kgroup_index(0); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B1[0]); + ++this->warp_tile_iterator_B_; + + int smem_write_stage_idx = Base::kStages - 1; + int smem_read_stage_idx = 0; + + warp_mma1.transform( + warp_transformed_frag_A1[0], + warp_transformed_frag_B1[0], + FragmentAScaler::apply( + warp_loaded_frag_A1[0], warp_loaded_frag_A1_scale[0]), + warp_loaded_frag_B1[0]); + + // tf32x3 kernels use staging accumulation. warp_mma uses a temporary + // accumulator and this temporary accumulator is added to the final + // accumulator once in every mainloop iteration. + plus plus_accum; + + FragmentC1 tmp_accum; + + if (platform::is_same< + typename Operator1::MathOperator, + arch::OpMultiplyAddFastF32>::value || + platform::is_same< + typename Operator1::MathOperator, + arch::OpMultiplyAddComplexFastF32>::value) { + tmp_accum.clear(); + } + + // + // Mainloop + // + + CUTLASS_PRAGMA_UNROLL + for (int gemm_k_iterations_1 = gemm_k_iterations_1_ - (Base::kStages - 1); + gemm_k_iterations_1 > (-Base::kStages + 1); + gemm_k_iterations_1--) { + // + // Loop over GEMM K dimension + // + + // Computes a warp-level GEMM on data held in shared memory + // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate + CUTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations1; + ++warp_mma_k) { + // Load warp-level tile from accumulator fragment (A) + // or shared memory (operand B) + this->warp_tile_iterator_B_.set_kgroup_index( + (warp_mma_k + 1) % Base::kWarpGemmIterations1); + // skip warp tile loading for the last kgroup (we are out of the buf) + if (gemm_k_iterations_1 > (-Base::kStages + 2) || + warp_mma_k < Base::kWarpGemmIterations1 - 1) { + warp_tile_iterator_A1_.load( + warp_loaded_frag_A1[(warp_mma_k + 1) % 2]); + warp_tile_iterator_A1_scale_.load( + warp_loaded_frag_A1_scale[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_B_.load( + warp_loaded_frag_B1[(warp_mma_k + 1) % 2]); + } + ++warp_tile_iterator_A1_; + ++warp_tile_iterator_A1_scale_; + ++this->warp_tile_iterator_B_; + + if (warp_mma_k > 0) + warp_mma1.transform( + warp_transformed_frag_A1[warp_mma_k % 2], + warp_transformed_frag_B1[warp_mma_k % 2], + FragmentAScaler::apply( + warp_loaded_frag_A1[warp_mma_k % 2], + warp_loaded_frag_A1_scale[warp_mma_k % 2]), + warp_loaded_frag_B1[warp_mma_k % 2]); + + if (platform::is_same< + typename Operator1::MathOperator, + arch::OpMultiplyAddFastF32>::value || + platform::is_same< + typename Operator1::MathOperator, + arch::OpMultiplyAddComplexFastF32>::value) { + warp_mma1( + tmp_accum, + warp_transformed_frag_A1[warp_mma_k % 2], + warp_transformed_frag_B1[warp_mma_k % 2], + tmp_accum); + + if (warp_mma_k == 0) { + accum = plus_accum(accum, tmp_accum); + tmp_accum.clear(); + } + } else { + warp_mma1( + accum, + warp_transformed_frag_A1[warp_mma_k % 2], + warp_transformed_frag_B1[warp_mma_k % 2], + accum); + } + + // Issue global->shared copies for the this stage + if (warp_mma_k < Base::kWarpGemmIterations1 - 1) { + int group_start_iteration_B1; + + group_start_iteration_B1 = warp_mma_k * Detail::kAccessesPerGroupB1; + + if (!kSmemContainsEntireB) { + copy_tiles_and_advance_1(iterator_B1, group_start_iteration_B1); + } + } + + if (warp_mma_k + 2 == Base::kWarpGemmIterations1) { + int group_start_iteration_B1; + group_start_iteration_B1 = + (warp_mma_k + 1) * Detail::kAccessesPerGroupB1; + + if (!kSmemContainsEntireB) { + copy_tiles_and_advance_1(iterator_B1, group_start_iteration_B1); + } + + // Inserts a memory fence between stages of cp.async instructions. + cutlass::arch::cp_async_fence(); + + // Waits until kStages-2 stages have committed. + arch::cp_async_wait(); + __syncthreads(); + + // Move to the next stage + iterator_B1.add_tile_offset({1, 0}); + + this->smem_iterator_B1_.add_tile_offset({1, 0}); + + // Add negative offsets to return iterators to the 'start' of the + // circular buffer in shared memory + if (!kSmemContainsEntireB) { + if (smem_write_stage_idx == (Base::kStages - 1)) { + this->smem_iterator_B1_.add_tile_offset({-Base::kStages, 0}); + smem_write_stage_idx = 0; + } else { + ++smem_write_stage_idx; + } + + if (smem_read_stage_idx == (Base::kStages - 1)) { + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy1::kPartitionsK * + Base::kWarpGemmIterations1, + 0}); + smem_read_stage_idx = 0; + } else { + ++smem_read_stage_idx; + } + } + + iterator_B1.set_residual_tile(gemm_k_iterations_1 == 2); + iterator_B1.clear_mask(gemm_k_iterations_1 == 1); + } + + // Do any conversions feeding the first stage at the end of the loop so + // we can start right away on mma instructions + if (warp_mma_k + 1 == Base::kWarpGemmIterations1) + warp_mma1.transform( + warp_transformed_frag_A1[(warp_mma_k + 1) % 2], + warp_transformed_frag_B1[(warp_mma_k + 1) % 2], + FragmentAScaler::apply( + warp_loaded_frag_A1[(warp_mma_k + 1) % 2], + warp_loaded_frag_A1_scale[(warp_mma_k + 1) % 2]), + warp_loaded_frag_B1[(warp_mma_k + 1) % 2]); + } + } + + if (platform::is_same< + typename Operator1::MathOperator, + arch::OpMultiplyAddFastF32>::value || + platform::is_same< + typename Operator1::MathOperator, + arch::OpMultiplyAddComplexFastF32>::value) { + accum = plus_accum(accum, tmp_accum); + } + } +}; + +// Converts a "regular" Mma into their counterpart from shared memory +template < + typename Mma_, + int kMaxK, + typename WarpIteratorA_, + /// whether or not to apply elementwise multiplication of operand A by + /// another matrix in shared memory before usage in A @ B + bool kScaleOperandA, + bool kTransposeA = false> +struct DefaultMmaFromSharedMemory; + +// Mma pipelined +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + typename WarpIteratorA_, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Transformation applied to A operand + typename TransformA_, + /// Transformation applied to B operand + typename TransformB_, + // Max MMA problem size K + int kMaxK, + /// whether or not to apply elementwise multiplication of operand A by + /// another matrix in shared memory before usage in A @ B + bool kScaleOperandA, + bool kTransposeA> +struct DefaultMmaFromSharedMemory< + MmaPipelined< + Shape_, + IteratorA_, + SmemIteratorA_, + IteratorB_, + SmemIteratorB_, + ElementC_, + LayoutC_, + Policy_, + TransformA_, + TransformB_>, + kMaxK, + WarpIteratorA_, + kScaleOperandA, + kTransposeA> { + using RegularMma = MmaPipelined< + Shape_, + IteratorA_, + SmemIteratorA_, + IteratorB_, + SmemIteratorB_, + ElementC_, + LayoutC_, + Policy_, + TransformA_, + TransformB_>; + + using WarpShape = typename Policy_::Operator::Shape; + using InstructionShape = typename Policy_::Operator::InstructionShape; + using ArchMmaOperator = typename Policy_::Operator; + + static constexpr bool kIsTransposedA = false; + using WarpIteratorA = WarpIteratorA_; + using IteratorB = + typename cutlass::transform::threadblock::MakeIteratorResidualLast< + IteratorB_>::Iterator; + + using Mma = typename cutlass::gemm::threadblock::MmaPipelinedFromSharedMemory< + Shape_, + WarpIteratorA, + kScaleOperandA, + kMaxK, + IteratorB, + SmemIteratorB_, + ElementC_, + LayoutC_, + Policy_>; +}; + +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + typename WarpIteratorA_, + /// Cache operation for operand A + cutlass::arch::CacheOperation::Kind CacheOpA, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Cache operation for operand B + cutlass::arch::CacheOperation::Kind CacheOpB, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear, + int kMaxK, + /// whether or not to apply elementwise multiplication of operand A by + /// another matrix in shared memory before usage in A @ B + bool kScaleOperandA, + bool kTransposeA> +struct DefaultMmaFromSharedMemory< + MmaMultistage< + Shape_, + IteratorA_, + SmemIteratorA_, + CacheOpA, + IteratorB_, + SmemIteratorB_, + CacheOpB, + ElementC_, + LayoutC_, + Policy_, + Stages, + SharedMemoryClear>, + kMaxK, + WarpIteratorA_, + kScaleOperandA, + kTransposeA> { + using RegularMma = MmaMultistage< + Shape_, + IteratorA_, + SmemIteratorA_, + CacheOpA, + IteratorB_, + SmemIteratorB_, + CacheOpB, + ElementC_, + LayoutC_, + Policy_, + Stages, + SharedMemoryClear>; + + using WarpShape = typename Policy_::Operator::Shape; + using InstructionShape = typename Policy_::Operator::InstructionShape; + using WarpIteratorTranspose = TransposeWarpIterator; + static constexpr bool kIsTransposedA = + WarpIteratorTranspose::kSupportsTranspose && kTransposeA; + using WarpIteratorA = typename platform::conditional< + kIsTransposedA, + typename WarpIteratorTranspose::Iterator, + WarpIteratorA_>::type; + + // Reduce the number of stages if we don't need that many + static int constexpr kStagesMax = + (kMaxK + int(Shape_::kK) - 1) / int(Shape_::kK); + static int constexpr kStages = cutlass::const_min(Stages, kStagesMax); + + using IteratorB = + typename cutlass::transform::threadblock::MakeIteratorResidualLast< + IteratorB_>::Iterator; + using Mma = + typename cutlass::gemm::threadblock::MmaMultistageFromSharedMemory< + Shape_, + WarpIteratorA, + kScaleOperandA, + IteratorB, + SmemIteratorB_, + RegularMma::kCacheOpB, + ElementC_, + LayoutC_, + Policy_, + kStages, + kMaxK>; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename IteratorC, + typename Operator, + typename scalar_t, + typename WarpShape_, + typename ThreadblockShape_> +struct B2bGemm; + +// Tensor Cores >= Sm75 specialization (Ampere ...) +template < /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Element type + typename Element_, + /// Layout of operand in memory + typename Layout_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions, concept: MatrixShape) + typename OpDelta_, + typename Operator, + typename scalar_t, + typename WarpShape_, + typename ThreadblockShape_> +struct B2bGemm< + cutlass::gemm::warp::MmaTensorOpAccumulatorTileIterator< + Shape_, + Element_, + Layout_, + InstructionShape_, + OpDelta_>, + Operator, + scalar_t, + WarpShape_, + ThreadblockShape_> { + using IteratorC = + typename cutlass::gemm::warp::MmaTensorOpAccumulatorTileIterator< + Shape_, + Element_, + Layout_, + InstructionShape_, + OpDelta_>; + using FragmentC = typename IteratorC::Fragment; + using InstructionShape = InstructionShape_; + using WarpShape = WarpShape_; + using ThreadblockShape = ThreadblockShape_; + using accum_t = Element_; + using lse_scalar_t = float; + + using SmemAccumulatorLayout = cutlass::layout::RowMajor; + + // Iterator to load accumulators (results of matmul in registers) + using FragmentIteratorAccumulator = + cutlass::epilogue::warp::FragmentIteratorTensorOp< + WarpShape, + InstructionShape, + accum_t, + typename Operator::Policy::Operator::FragmentC, + cutlass::layout::RowMajor>; + + // Iterator to store to shared-memory + using SmemIteratorD0 = typename cutlass::epilogue::warp::TileIteratorTensorOp< + WarpShape, + InstructionShape, + scalar_t, // accum_t, + SmemAccumulatorLayout>; + using AccumulatorSharedStorage = + cutlass::gemm::threadblock::AccumulatorSharedStorage< + ThreadblockShape, + typename SmemIteratorD0::Element, + typename SmemIteratorD0::TensorLayout, + typename SmemIteratorD0::Padding>; + // We need to provide an operation for the epilogue. Let's create an + // operation that does nothing (ScaleType::Nothing), just converts + // from accum_t (float) -> scalar_t (can be half) + using OutputOpNoOp = cutlass::epilogue::thread::LinearCombination< + typename SmemIteratorD0::Element, // ElementOutput + FragmentIteratorAccumulator::Fragment::kElements, + accum_t, // ElementAccumulator + typename SmemIteratorD0::Element, // ElementCompute + cutlass::epilogue::thread::ScaleType::Nothing>; + using Epilogue = cutlass::epilogue::threadblock::EpilogueSmemAccumulator< + SmemIteratorD0, + FragmentIteratorAccumulator, + SmemIteratorD0, // ScaleBiasIterator - not used + OutputOpNoOp>; + + // Epilogue 2: with LSE (for backwards pass) + using EpilogueOpApplyExp = cutlass::epilogue::thread::ApplyExp< + scalar_t, // ElementOutput_ + lse_scalar_t, // ElementLSE_ + accum_t, // ElementAccumulator_ + accum_t, // ElementCompute_ + 128 / cutlass::sizeof_bits::value>; + using EpilogueWithExp = + cutlass::epilogue::threadblock::EpilogueSmemAccumulator< + SmemIteratorD0, // ScaleBiasIterator - not used + FragmentIteratorAccumulator, + SmemIteratorD0, // ScaleBiasIterator - not used + EpilogueOpApplyExp>; + + static void CUTLASS_DEVICE accumToSmem( + AccumulatorSharedStorage& shared_storage, + FragmentC const& accum, + int lane_id, + cutlass::MatrixCoord const& tile_coords) { + SmemIteratorD0 smem_iterator_attn(shared_storage.accum_ref(), lane_id); + smem_iterator_attn.add_tile_offset( + tile_coords * + cutlass::MatrixCoord{ + SmemIteratorD0::TileIterations::kRow, + SmemIteratorD0::TileIterations::kColumn}); + Epilogue epilogue; + epilogue(OutputOpNoOp({}), smem_iterator_attn, accum); + } + + // NOTE(alih): we only apply the exp operator here; + // LSE is pre-fetched into shared memory, and + // subtracted when we check the NA mask and set + // invalid weights to -info. + // Refer to kernel_backward.h for more details. + static void CUTLASS_DEVICE accumApplyExpToSmem( + AccumulatorSharedStorage& shared_storage, + FragmentC& accum, + int thread_id, + int warp_id, + int lane_id, + cutlass::MatrixCoord const& tile_coords) { + SmemIteratorD0 smem_iterator_attn(shared_storage.accum_ref(), lane_id); + smem_iterator_attn.add_tile_offset( + tile_coords * + cutlass::MatrixCoord{ + SmemIteratorD0::TileIterations::kRow, + SmemIteratorD0::TileIterations::kColumn}); + EpilogueWithExp epilogue; + EpilogueOpApplyExp exp({}); + epilogue(exp, smem_iterator_attn, accum); + } +}; + +// Volta Specialization +// only supported for f16 +template +struct B2bGemm< + cutlass::gemm::warp::MmaVoltaTensorOpAccumulatorTileIterator< + cutlass::MatrixShape<32, 32>, + float, + cutlass::layout::RowMajor, + cutlass::gemm::GemmShape<16, 16, 4>, + cutlass::MatrixShape<1, 1>>, + Operator, + cutlass::half_t, + WarpShape_, + ThreadblockShape_> { + using IteratorC = + cutlass::gemm::warp::MmaVoltaTensorOpAccumulatorTileIterator< + cutlass::MatrixShape<32, 32>, + float, + cutlass::layout::RowMajor, + cutlass::gemm::GemmShape<16, 16, 4>, + cutlass::MatrixShape<1, 1>>; + using scalar_t = cutlass::half_t; + using accum_t = IteratorC::Element; + using WarpShape = WarpShape_; + using ThreadblockShape = ThreadblockShape_; + using FragmentC = IteratorC::Fragment; + using lse_scalar_t = float; + + // Storage in shared-memory for Q.Kt + using SmemAccumulatorLayout = + cutlass::layout::RowMajorVoltaTensorOpMultiplicandCrosswise<16, 32>; + using AccumulatorSharedStorage = + cutlass::gemm::threadblock::AccumulatorSharedStorage< + ThreadblockShape, + scalar_t, + SmemAccumulatorLayout, + cutlass::MatrixShape<0, 0> // Padding + >; + using TensorRef = cutlass::TensorRef; + using Policy = typename IteratorC::Policy; + using Element = accum_t; + // Those are MmaVoltaTensorOpAccumulatorTileIterator private fields + // Let's copy their values + static int const kElementsPerPartial = 4; + using EleShapePerPatial = typename cutlass::platform::conditional< + cutlass::platform::is_same::value, + cutlass::MatrixShape<2, 2>, + cutlass::MatrixShape<1, 4>>::type; + static int const kElementsPerMma = 8; + static int const kAccumulatorPatials = 2; + using QuadShapePerPatialMma = cutlass::MatrixShape<4, 4>; + + static void CUTLASS_DEVICE accumToSmem( + AccumulatorSharedStorage& shared_storage, + FragmentC const& accum, + int lane_id, + cutlass::MatrixCoord const& tile_coords) { + // ctor - from MmaVoltaTensorOpAccumulatorTileIterator + TensorRef ref_(shared_storage.accum_ref()); + int quad = (lane_id >> 2); + int lane_in_quad = (lane_id & 3); + int accum_m, accum_n; + + if (cutlass::platform::is_same::value) { + // (quad[2],quad[0])+lane_in_quad[0] + accum_m = (((quad & 0x4) >> 1) + (quad & 0x1)) * 8 + (lane_in_quad & 1); + // (quad[1])+lane_in_quad[1] + accum_n = + ((quad >> 1) & 0x1) * kElementsPerPartial * kAccumulatorPatials + + (lane_in_quad & 2); + } else { + accum_m = (((quad & 0x4) >> 1) + (quad & 0x1)) * 8 + + lane_in_quad; // (quad[2],quad[0]) + accum_n = ((quad >> 1) & 0x1) * kElementsPerPartial * kAccumulatorPatials; + } + cutlass::MatrixCoord lane_offset(accum_m, accum_n); + + // Tile offset + ref_.add_coord_offset( + tile_coords * + cutlass::MatrixCoord( + {IteratorC::Shape::kRow, IteratorC::Shape::kColumn})); + + using AccessType = cutlass::Array; + + // store - from MmaVoltaTensorOpAccumulatorTileIterator + CUTLASS_PRAGMA_UNROLL + for (int tile_n = 0; tile_n < Policy::TileIterations::kColumn; ++tile_n) { + CUTLASS_PRAGMA_UNROLL + for (int tile_m = 0; tile_m < Policy::TileIterations::kRow; ++tile_m) { + CUTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + CUTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + int mma_accum_start = + (((tile_n * Policy::TileIterations::kRow + tile_m) * + Policy::MmaIterations::kColumn + + mma_n) * + Policy::MmaIterations::kRow + + mma_m) * + kElementsPerMma; + + CUTLASS_PRAGMA_UNROLL + for (int p = 0; p < kAccumulatorPatials; ++p) { + CUTLASS_PRAGMA_UNROLL + for (int m = 0; m < EleShapePerPatial::kRow; ++m) { + int accum_m = tile_m * Policy::InterleavedTile::kRow + + mma_m * QuadShapePerPatialMma::kRow + m * 2; + int accum_n = tile_n * Policy::InterleavedTile::kColumn + + mma_n * QuadShapePerPatialMma::kColumn + + p * Policy::InterleavedTile::kColumn / 2; + int r = (accum_m + lane_offset.row()); + AccessType to_store; + CUTLASS_PRAGMA_UNROLL + for (int n = 0; n < EleShapePerPatial::kColumn; ++n) { + int idx = mma_accum_start + p * kElementsPerPartial + + m * EleShapePerPatial::kColumn + n; + int c = (accum_n + n + lane_offset.column()); + to_store[n] = scalar_t(accum[idx]); + } + int c = (accum_n + lane_offset.column()); + assert(r < 32); + assert(c < 32); + *reinterpret_cast( + ref_.data() + ref_.offset({r, c})) = to_store; + } + } + } + } + } + } + } + + // NOTE(alih): we only apply the exp operator here; + // LSE is pre-fetched into shared memory, and + // subtracted when we check the NA mask and set + // invalid weights to -info. + // Refer to kernel_backward.h for more details. + static void CUTLASS_DEVICE accumApplyExpToSmem( + AccumulatorSharedStorage& shared_storage, + typename IteratorC::Fragment& accum, + int thread_id, + int warp_id, + int lane_id, + cutlass::MatrixCoord const& tile_coords) { + // Non-optimized way to apply exp to registers + // NOTE: accum is attn.T + // TODO: Optimize for each architecture + static constexpr int WarpSize = 32; + using AccumLambdaIterator = + typename DefaultMmaAccumLambdaIterator:: + Iterator; + auto lane_offset = + AccumLambdaIterator::get_lane_offset(lane_id, warp_id, tile_coords); + + AccumLambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) {}, + [&](int accum_m, int accum_n, int idx) { + accum[idx] = expf(accum[idx]); + }, + [&](int accum_m) {}); + accumToSmem(shared_storage, accum, lane_id, tile_coords); + } +}; + +// Simt Specialization +// for f32 on Sm70-Sm75 and f16/f32 below + +template < + typename Operator, + typename OperatorPolicy, + typename scalar_t, + typename WarpShape_, + typename ThreadblockShape_> +struct B2bGemm< + cutlass::gemm::warp::MmaSimtTileIterator< + cutlass::MatrixShape<32, 32>, + cutlass::gemm::Operand::kC, + float, + cutlass::layout::RowMajor, + OperatorPolicy, + 1, + 1>, + Operator, + scalar_t, + WarpShape_, + ThreadblockShape_> { + using IteratorC = cutlass::gemm::warp::MmaSimtTileIterator< + cutlass::MatrixShape<32, 32>, + cutlass::gemm::Operand::kC, + float, + cutlass::layout::RowMajor, + OperatorPolicy, + 1, + 1>; + using accum_t = typename IteratorC::Element; + using WarpShape = WarpShape_; + using ThreadblockShape = ThreadblockShape_; + using FragmentC = typename IteratorC::Fragment; + using lse_scalar_t = float; + + // Storage in shared-memory for Q.Kt + using AccumulatorSharedStorage = + cutlass::gemm::threadblock::AccumulatorSharedStorage< + ThreadblockShape, + scalar_t, + cutlass::layout::ColumnMajor, + cutlass::MatrixShape<0, 0> // Padding + >; + + static void CUTLASS_DEVICE accumToSmem( + AccumulatorSharedStorage& shared_storage, + FragmentC const& accum, + int lane_id, + cutlass::MatrixCoord const& tile_coords) { + using Policy = typename IteratorC::Policy; + using Element = typename IteratorC::Element; + using Iterations = typename IteratorC::Iterations; + using Delta = typename IteratorC::Delta; + + auto ref_ = shared_storage.accum_ref(); + // ctor - MmaSimtTileIterator + // compute offset based on thread ID and lane layout + typename Policy::LaneLayout lane_layout = Policy::get_lane_layout(); + + MatrixCoord lane_offset = lane_layout.inverse(lane_id) * + MatrixCoord(Policy::LaneMmaShape::kM, Policy::LaneMmaShape::kN); + + ref_.add_coord_offset(lane_offset); + + // Tile offset + ref_.add_coord_offset( + tile_coords * + cutlass::MatrixCoord( + {IteratorC::Shape::kRow, IteratorC::Shape::kColumn})); + + // store - MmaSimtTileIterator + CUTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Iterations::kColumn; ++mma_n) { + CUTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::LaneMmaShape::kN; ++n) { + CUTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Iterations::kRow; ++mma_m) { + CUTLASS_PRAGMA_UNROLL + for (int m = 0; m < Policy::LaneMmaShape::kM; ++m) { + int r = + Policy::LaneMmaShape::kM * (mma_m * Policy::WarpShape::kRow) + + m; + int c = mma_n * Delta::kColumn + n; + int idx = n + + Policy::LaneMmaShape::kN * + (mma_n + + Iterations::kColumn * + (m + mma_m * Policy::LaneMmaShape::kM)); + ref_.at({r, c}) = scalar_t(accum[idx]); + } + } + } + } + } + + // NOTE(alih): we only apply the exp operator here; + // LSE is pre-fetched into shared memory, and + // subtracted when we check the NA mask and set + // invalid weights to -info. + // Refer to kernel_backward.h for more details. + static void CUTLASS_DEVICE accumApplyExpToSmem( + AccumulatorSharedStorage& shared_storage, + typename IteratorC::Fragment& accum, + int thread_id, + int warp_id, + int lane_id, + cutlass::MatrixCoord const& tile_coords) { + // Non-optimized way to apply exp to registers + // NOTE: accum is attn.T + // TODO: Optimize for each architecture + static constexpr int WarpSize = 32; + using AccumLambdaIterator = + typename DefaultMmaAccumLambdaIterator:: + Iterator; + auto lane_offset = + AccumLambdaIterator::get_lane_offset(lane_id, warp_id, tile_coords); + + AccumLambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) {}, + [&](int accum_m, int accum_n, int idx) { + accum[idx] = expf(accum[idx]); + }, + [&](int accum_m) {}); + accumToSmem(shared_storage, accum, lane_id, tile_coords); + } +}; + +} // namespace threadblock +} // namespace gemm +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fmha/gemm_kernel_utils.h b/natten/csrc/include/natten/cuda/fmha/gemm_kernel_utils.h new file mode 100644 index 00000000..be00a5c6 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha/gemm_kernel_utils.h @@ -0,0 +1,240 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ +#pragma once + +#include "cutlass/arch/mma.h" + +//////////////////////////////////////////////////////////////////////////////// +// Some helper functions +//////////////////////////////////////////////////////////////////////////////// +#define DISPATCH_TYPES(tensor, func) \ + { \ + if (query.scalar_type() == at::ScalarType::Float) { \ + using scalar_t = float; \ + func(); \ + } else if (query.scalar_type() == at::ScalarType::Half) { \ + using scalar_t = cutlass::half_t; \ + func(); \ + } else if (query.scalar_type() == at::ScalarType::BFloat16) { \ + using scalar_t = cutlass::bfloat16_t; \ + func(); \ + } else { \ + XFORMERS_CHECK(false, "Only fp32, half & bf16 supported at the moment"); \ + } \ + } + +#define DISPATCH_BOOL(BOOL_V, BOOL_NAME, F) \ + { \ + if (BOOL_V) { \ + constexpr bool BOOL_NAME = true; \ + F(); \ + } else { \ + constexpr bool BOOL_NAME = false; \ + F(); \ + } \ + } +#define DISPATCH_ARCHTAG(CC, func) \ + { \ + if (CC >= 80) { \ + using ArchTag = cutlass::arch::Sm80; \ + func(); \ + } else if (CC >= 75) { \ + using ArchTag = cutlass::arch::Sm75; \ + func(); \ + } else if (CC >= 70) { \ + using ArchTag = cutlass::arch::Sm70; \ + func(); \ + } else if (CC >= 50) { \ + using ArchTag = cutlass::arch::Sm50; \ + func(); \ + } else { \ + XFORMERS_CHECK( \ + false, \ + "Your device is too old. We require compute capability >= 50"); \ + } \ + } + +#define CHECK_NOSPARSE_CONTIGUOUS_CUDA(TENSOR) \ + XFORMERS_CHECK(TENSOR.is_cuda(), #TENSOR " must be a CUDA tensor"); \ + XFORMERS_CHECK(!TENSOR.is_sparse(), #TENSOR " must be a dense tensor"); \ + XFORMERS_CHECK(TENSOR.is_contiguous()); + +#define CHECK_NOSPARSE_LASTCONTIGUOUS_CUDA(TENSOR) \ + XFORMERS_CHECK(TENSOR.is_cuda(), #TENSOR " must be a CUDA tensor"); \ + XFORMERS_CHECK(!TENSOR.is_sparse(), #TENSOR " must be a dense tensor"); \ + XFORMERS_CHECK( \ + TENSOR.stride(-1) == 1, #TENSOR ": last dimension must be contiguous"); + +#ifdef TORCH_CHECK +#define CHECK_ALIGNED_PTR(PTR, ALIGNMENT) \ + XFORMERS_CHECK( \ + uint64_t(PTR) % ALIGNMENT == 0, #PTR " is not correctly aligned") +#define XFORMERS_CHECK TORCH_CHECK +#elif defined(__CUDACC_RTC__) +#define CHECK_ALIGNED_PTR(PTR, ALIGNMENT) \ + if (!(uint64_t(PTR) % ALIGNMENT == 0)) { \ + return false; \ + } +#define XFORMERS_CHECK(COND, ERR) \ + if (!(COND)) { \ + return false; \ + } +#else +#include +#define CHECK_ALIGNED_PTR(PTR, ALIGNMENT) \ + if (!(uint64_t(PTR) % ALIGNMENT == 0)) { \ + std::cerr << #PTR " is not correctly aligned\n"; \ + return false; \ + } + +template +void print_check_args(const Args&... args) { + (std::cerr << ... << args) << std::endl; +} + +#define XFORMERS_CHECK(COND, ...) \ + if (!(COND)) { \ + std::cerr << "'" #COND "' failed: "; \ + print_check_args(__VA_ARGS__); \ + return false; \ + } +#endif + +#define ASSIGN_CHECK_OVERFLOW(A, B) \ + { \ + A = B; \ + XFORMERS_CHECK( \ + B < std::numeric_limits::max(), #B " overflows"); \ + } + +namespace gemm_kernel_utils { + +template +constexpr CUTLASS_HOST_DEVICE integer ceil_div(integer n, integer m) { + return (n + m - 1) / m; +} + +template +constexpr CUTLASS_HOST_DEVICE integer align_up(integer n, integer m) { + return ((n + m - 1) / m) * m; +} + +//////////////////////////////////////////////////////////////////////////////// +// Determine the type of GEMM we do (TensorCores or not, Shapes ...) +// TODO: Maybe we could rely on Cutlass's DefaultGemm templates +//////////////////////////////////////////////////////////////////////////////// + +// Fallback to Simt (FMA on cuda cores) if not in a special case below +template +struct DefaultGemmType { + static constexpr int ThreadK = 8; + static constexpr int WarpK = 8; + static constexpr int kMinimumAlignment = 1; + using InstructionShape = cutlass::gemm::GemmShape<1, 1, 1>; + using OpClass = cutlass::arch::OpClassSimt; + using Operator = cutlass::arch::OpMultiplyAdd; +}; + +// Specialization for tensorcores with f32 +template +struct DefaultGemmType< + ArchTag, + float, + typename cutlass::platform::enable_if< + ArchTag::kMinComputeCapability >= 80>::type> { + static constexpr int ThreadK = 32; + static constexpr int WarpK = 32; + static constexpr int kMinimumAlignment = 4; + using OpClass = cutlass::arch::OpClassTensorOp; + using InstructionShape = cutlass::gemm::GemmShape<16, 8, 8>; + using Operator = cutlass::arch::OpMultiplyAddFastF32; +}; + +// Specialization for tensorcores with f16/bf16 - Sm75+ +template +struct DefaultGemmType< + ArchTag, + scalar_t, + typename cutlass::platform::enable_if< + ArchTag::kMinComputeCapability >= 75 && + cutlass::sizeof_bits::value == 16>::type> { + static constexpr int ThreadK = 32; + static constexpr int WarpK = 32; + static constexpr int kMinimumAlignment = 4; + using OpClass = cutlass::arch::OpClassTensorOp; + using InstructionShape = cutlass::gemm::GemmShape<16, 8, 8>; + using Operator = cutlass::arch::OpMultiplyAdd; +}; + +// Specialization for tensorcores with f16 - Volta +template <> +struct DefaultGemmType { + static constexpr int ThreadK = 32; + static constexpr int WarpK = 32; + static constexpr int kMinimumAlignment = 2; + using OpClass = cutlass::arch::OpClassTensorOp; + using InstructionShape = cutlass::gemm::GemmShape<8, 8, 4>; + using Operator = cutlass::arch::OpMultiplyAdd; +}; + +// Enables to do +// `auto x = kCondition ? fa(arg) : fb(arg)` +// when `fa` and `fb` have different types +template +struct call_conditional; + +template +struct call_conditional { + template + static CUTLASS_HOST_DEVICE auto apply(TA ta, TB tb, Arg arg) + -> decltype(ta(arg)) { + return ta(arg); + } +}; + +template +struct call_conditional { + template + static CUTLASS_HOST_DEVICE auto apply(TA ta, TB tb, Arg arg) + -> decltype(tb(arg)) { + return tb(arg); + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// Mark a variable as warp-uniform - enables some compiler optimizations +// The cheapest way to do it is just to broadcast it from lane 0 +//////////////////////////////////////////////////////////////////////////////// + +template +CUTLASS_DEVICE T warp_uniform(T value) { + struct { + union { + T value; + uint32_t asInt; + }; + } p; + p.value = value; + p.asInt = __shfl_sync(0xffffffff, (unsigned)p.asInt, 0); + return p.value; +} + +template +CUTLASS_DEVICE T* warp_uniform(T* ptr) { + struct { + union { + T* ptr; + uint32_t asInt[2]; + }; + } p; + p.ptr = ptr; + p.asInt[0] = warp_uniform(p.asInt[0]); + p.asInt[1] = warp_uniform(p.asInt[1]); + return p.ptr; +} +} // namespace gemm_kernel_utils diff --git a/natten/csrc/include/natten/cuda/fmha/iterators/default_warp_iterator_from_smem.h b/natten/csrc/include/natten/cuda/fmha/iterators/default_warp_iterator_from_smem.h new file mode 100644 index 00000000..d75db8f8 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha/iterators/default_warp_iterator_from_smem.h @@ -0,0 +1,149 @@ +/* + * Copied from xFormers (https://github.com/facebookresearch/xformers/) + * + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + */ +/*************************************************************************************************** + * Copyright (c) 2017 - 2026 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 Instanciates the right WarpIterator to read from shared memory + The class `DefaultWarpIteratorAFromSharedMemory` is useful when reading + data dumped with `B2bGemm::accumToSmem`. +*/ + +#pragma once + +#include +#include +#include + +#include + +namespace cutlass { +namespace gemm { +namespace threadblock { + +template < + typename WarpShape, + typename InstructionShape, + typename RegularWarpIterator, + typename Policy, + typename Enable = void> +struct DefaultWarpIteratorAFromSharedMemory {}; + +// TensorOp - Ampere half +template +struct DefaultWarpIteratorAFromSharedMemory< + cutlass::gemm::GemmShape<32, 32, 32>, + cutlass::gemm::GemmShape<16, 8, kInstrK>, + RegularWarpIterator, + Policy, + typename platform::enable_if<( + sizeof_bits::value == 16 && + Policy::Operator::Policy::OpDelta::kRow == 1)>::type> { + using OpDelta = typename Policy::Operator::Policy::OpDelta; + using WarpShape = cutlass::MatrixShape<32, 32>; + using InstructionShape = cutlass::gemm::GemmShape<16, 8, kInstrK>; + + using WarpIterator = cutlass::gemm::warp::WarpIteratorFromSmem< + cutlass::gemm::Operand::kA, + typename RegularWarpIterator::Element, + cutlass::MatrixShape>; +}; + +// TensorOp - Ampere f32 +template +struct DefaultWarpIteratorAFromSharedMemory< + WarpShape, + cutlass::gemm::GemmShape<16, 8, 8>, + RegularWarpIterator, + Policy, + typename platform::enable_if<( + sizeof_bits::value != 16 || + Policy::Operator::Policy::OpDelta::kRow != 1)>::type> { + using InstructionShape = cutlass::gemm::GemmShape<16, 8, 8>; + static constexpr auto kWarpSize = 32; + using OpDelta = typename Policy::Operator::Policy::OpDelta; + + using WarpIterator = + cutlass::gemm::warp::MmaTensorOpMultiplicandTileAccessIterator< + cutlass::MatrixShape, + cutlass::gemm::Operand::kA, + typename RegularWarpIterator::Element, + cutlass::layout::RowMajor, + cutlass::MatrixShape, + OpDelta::kRow, + kWarpSize>; +}; + +// TensorOp - Volta +template +struct DefaultWarpIteratorAFromSharedMemory< + WarpShape, + cutlass::gemm::GemmShape<16, 16, 4>, + RegularWarpIterator, + Policy> { + using InstructionShape = cutlass::gemm::GemmShape<16, 16, 4>; + static constexpr auto kWarpSize = 32; + using OpDelta = typename Policy::Operator::Policy::OpDelta; + + using WarpIterator = + cutlass::gemm::warp::MmaVoltaTensorOpMultiplicandTileIterator< + cutlass::MatrixShape<32, 32>, // MatrixShape, + cutlass::gemm::Operand::kA, + typename RegularWarpIterator::Element, + cutlass::layout::RowMajorVoltaTensorOpMultiplicandCrosswise<16, 32>, + cutlass::MatrixShape<16, 4>, + OpDelta::kRow, + kWarpSize>; +}; + +// Simt +template +struct DefaultWarpIteratorAFromSharedMemory< + WarpShape, + cutlass::gemm::GemmShape<1, 1, 1>, + RegularWarpIterator, + Policy> { + using InstructionShape = cutlass::gemm::GemmShape<1, 1, 1>; + static constexpr auto kWarpSize = 32; + + // We just use the same iterator, as we reproduced the same shared-memory + // schema. Just modify it to handle non-complete tiles. + using WarpIterator = RegularWarpIterator; +}; + +} // namespace threadblock +} // namespace gemm +} // namespace cutlass diff --git a/natten/csrc/include/natten/cuda/fmha/iterators/epilogue_predicated_tile_iterator.h b/natten/csrc/include/natten/cuda/fmha/iterators/epilogue_predicated_tile_iterator.h new file mode 100644 index 00000000..1ece8b9f --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha/iterators/epilogue_predicated_tile_iterator.h @@ -0,0 +1,760 @@ +/* + * Copied from xFormers (https://github.com/facebookresearch/xformers/) and + * edited. + * + * NOTE: this is only used in the backward pass kernel. + * + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + */ +/*************************************************************************************************** + * Copyright (c) 2017 - 2026 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 Epilogue iterator that supports prefetching + + Mostly copied from "cutlass/epilogue/threadblock/predicated_tile_iterator.h" +*/ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { + +//////////////////////////////////////////////////////////////////////////////// + +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator used to load and store output tile from global memory in +/// epilogue. +/// +/// Satisfies: ReadableTileIterator | PredicatedTileIterator | +/// ForwardTileIterator +/// +template < + typename ThreadMap_, ///< Thread map (conept: OutputTileThreadMap) + typename Element_, ///< Element data type + bool ScatterD = false, ///< Scatter D operand or not + bool UseCUDAStore = false> +class PredicatedTileIteratorPrefetch { + public: + using ThreadMap = ThreadMap_; + using Shape = typename ThreadMap::Shape; + + using Element = Element_; + + using Layout = layout::RowMajor; + using TensorRef = TensorRef; + using ConstTensorRef = typename TensorRef::ConstTensorRef; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using TensorCoord = MatrixCoord; + + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + static int const kThreads = ThreadMap::kThreads; + static int const kIterations = ThreadMap::Count::kTile; + + static_assert( + ThreadMap::Iterations::kRow > 0, + "ThreadMap::Iterations::kRow must be > 0"); + static_assert( + ThreadMap::Iterations::kGroup > 0, + "ThreadMap::Iterations::kGroup must be > 0"); + static_assert( + ThreadMap::Iterations::kCluster > 0, + "ThreadMap::Iterations::kCluster must be > 0"); + static_assert( + ThreadMap::Iterations::kColumn > 0, + "ThreadMap::Iterations::kColumn must be > 0"); + + /// Fragment object + using Fragment = Array< + Element, + ThreadMap::Iterations::kColumn * ThreadMap::Iterations::kRow * + ThreadMap::Iterations::kGroup * ThreadMap::Iterations::kCluster * + ThreadMap::kElementsPerAccess>; + + /// Memory access size + using AccessType = AlignedArray; + + // + // Parameters struct + // + + /// Uses a non-template class + struct Params : PredicatedTileIteratorParams { + using Base = PredicatedTileIteratorParams; + + CUTLASS_HOST_DEVICE + Params() {} + + CUTLASS_HOST_DEVICE + Params(Layout const& layout) + : PredicatedTileIteratorParams( + layout.stride(0) * int(sizeof(AccessType)) / kElementsPerAccess, + make_OutputTileThreadMapDesc()) {} + + CUTLASS_HOST_DEVICE + Params(Base const& base) : Base(base) {} + }; + + /// Mask object + struct Mask { + static int const kCount = ThreadMap::Iterations::kColumn; + + /// Predicate state + bool predicates[kCount]; + + // + // Mask + // + CUTLASS_HOST_DEVICE + Mask() { + enable(); + } + + ///< Efficiently disables all accesses guarded by mask + CUTLASS_HOST_DEVICE void clear() { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + predicates[i] = false; + } + } + + ///< CUTLASS_HOST_DEVICE enables all accesses guarded by mask + CUTLASS_DEVICE void enable() { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + predicates[i] = true; + } + } + }; + + private: + // + // Data members + // + + /// Parameters structure containing reference and precomputed state. + PredicatedTileIteratorParams params_; + + /// Byte-level pointer + uint8_t* byte_pointer_; + + /// Array of boolean values to contain steady-state predicates + Mask mask_; + + /// Extent of the matrix tile in rows + Index extent_row_; + + /// Extent of the matrix tile in rows + Index extent_column_; + + /// A thread's starting row position (assuming steady-state predicates have + /// been computed) + Index thread_start_row_; + + /// A thread's starting column + Index thread_start_column_; + + /// Internal state counter + int state_[3]; + + /// Scatter indices + int const* indices_; + + // + // Static asserts about internal strides + // + + static_assert(sizeof(extent_row_) == 4, "Expected 32b extents"); + static_assert(sizeof(thread_start_row_) == 4, "Expected 32b extents"); + static_assert( + sizeof(PredicatedTileIteratorParams::stride) == 8, + "Expected 64b strides"); + + private: + // + // Methods + // + + public: + // + // Methods + // + + /// Constructor + CUTLASS_DEVICE + PredicatedTileIteratorPrefetch( + PredicatedTileIteratorParams const& params, + Element* pointer, + TensorCoord extent, + int thread_idx, + TensorCoord threadblock_offset = TensorCoord(), + int const* indices = nullptr) + : params_(params), indices_(indices) { + TensorCoord thread_offset = + ThreadMap::initial_offset(thread_idx) + threadblock_offset; + + extent_row_ = extent.row(); + extent_column_ = extent.column(); + + thread_start_row_ = thread_offset.row(); + thread_start_column_ = thread_offset.column(); + + // Initialize predicates + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kColumn; ++c) { + mask_.predicates[c] = + ((thread_offset.column() + ThreadMap::Delta::kColumn * c) < + extent.column()); + } + + // Null pointer performs no accesses + if (!pointer) { + mask_.clear(); + } + + if (ScatterD && !indices) { + mask_.clear(); + } + + // Initialize pointer + byte_pointer_ = reinterpret_cast(pointer) + + LongIndex(thread_offset.row()) * LongIndex(params_.stride) + + LongIndex(thread_offset.column()) * sizeof(AccessType) / + kElementsPerAccess; + + if (ScatterD) { + byte_pointer_ = reinterpret_cast(pointer) + + LongIndex(thread_offset.column()) * sizeof(AccessType) / + kElementsPerAccess; + } + + // Initialize internal state counter + state_[0] = state_[1] = state_[2] = 0; + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + byte_pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + CUTLASS_DEVICE + void prefetch_all() { + CUTLASS_PRAGMA_UNROLL + for (int iter = 0; iter < kIterations; ++iter) { + prefetch(); + ++(*this); + } + } + + CUTLASS_DEVICE + void prefetch() { + uint8_t* byte_pointer = byte_pointer_; + + CUTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; + ++cluster) { + CUTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + CUTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + int row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + + AccessType* memory_pointer = + reinterpret_cast(byte_pointer); + + CUTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; + ++column) { + // on windows using unsigned long here gives the error + // error: asm operand type size(4) does not match + // type/size implied by constraint 'l' + uint64_t addr = (uint64_t)((void*)&memory_pointer + [column * ThreadMap::Delta::kColumn / + kElementsPerAccess]); + asm volatile("prefetch.global.L1 [ %1 ];" : "=l"(addr) : "l"(addr)); + } + + if (row + 1 < ThreadMap::Iterations::kRow) { + if (!ScatterD) { + byte_pointer += params_.increment_row; + } + } + } + + if (group + 1 < ThreadMap::Iterations::kGroup) { + byte_pointer += params_.increment_group; + } + } + + if (cluster + 1 < ThreadMap::Iterations::kCluster) { + byte_pointer += params_.increment_cluster; + } + } + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_byte_offset(Fragment& frag, int64_t byte_offset) const { + uint8_t* byte_pointer = byte_pointer_; + AccessType* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; + ++cluster) { + CUTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + CUTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + int frag_row_idx = + (row + + ThreadMap::Iterations::kRow * + (group + ThreadMap::Iterations::kGroup * cluster)); + + int row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + + bool row_guard = ((row_offset + thread_start_row_) < extent_row_); + + AccessType* memory_pointer = + reinterpret_cast(byte_pointer + byte_offset); + + if (ScatterD && row_guard) { + assert(indices_); + + memory_pointer = reinterpret_cast( + byte_pointer + byte_offset + + LongIndex(indices_[row_offset + thread_start_row_]) * + LongIndex(params_.stride)); + } + + CUTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; + ++column) { + bool guard = row_guard && mask_.predicates[column]; + + cutlass::arch::global_load( + frag_ptr + [frag_row_idx * ThreadMap::Iterations::kColumn + column], + (void*)&memory_pointer + [column * ThreadMap::Delta::kColumn / kElementsPerAccess], + guard); + } + + if (row + 1 < ThreadMap::Iterations::kRow) { + if (!ScatterD) { + byte_pointer += params_.increment_row; + } + } + } + + if (group + 1 < ThreadMap::Iterations::kGroup) { + byte_pointer += params_.increment_group; + } + } + + if (cluster + 1 < ThreadMap::Iterations::kCluster) { + byte_pointer += params_.increment_cluster; + } + } + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) const { + load_with_byte_offset(frag, 0); + } + + /// Stores a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, int64_t byte_offset) const { + uint8_t* byte_pointer = byte_pointer_; + AccessType const* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; + ++cluster) { + CUTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + CUTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + int frag_row_idx = + (row + + ThreadMap::Iterations::kRow * + (group + ThreadMap::Iterations::kGroup * cluster)); + + int row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + + bool row_guard = ((row_offset + thread_start_row_) < extent_row_); + + AccessType* memory_pointer = + reinterpret_cast(byte_pointer + byte_offset); + + if (ScatterD && row_guard) { + assert(indices_); + + memory_pointer = reinterpret_cast( + byte_pointer + byte_offset + + LongIndex(indices_[row_offset + thread_start_row_]) * + LongIndex(params_.stride)); + } + + CUTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; + ++column) { + bool guard = row_guard && mask_.predicates[column]; + + if (UseCUDAStore) { + if (guard) { + memory_pointer + [column * ThreadMap::Delta::kColumn / kElementsPerAccess] = + frag_ptr + [frag_row_idx * ThreadMap::Iterations::kColumn + + column]; + } + } else { + cutlass::arch::global_store( + frag_ptr + [frag_row_idx * ThreadMap::Iterations::kColumn + column], + (void*)&memory_pointer + [column * ThreadMap::Delta::kColumn / kElementsPerAccess], + guard); + } + } + + if (row + 1 < ThreadMap::Iterations::kRow) { + if (!ScatterD) { + byte_pointer += params_.increment_row; + } + } + } + + if (group + 1 < ThreadMap::Iterations::kGroup) { + byte_pointer += params_.increment_group; + } + } + + if (cluster + 1 < ThreadMap::Iterations::kCluster) { + byte_pointer += params_.increment_cluster; + } + } + } + + /// Stores a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) const { + store_with_byte_offset(frag, 0); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void downsample_load_with_byte_offset( + Fragment& frag, + int64_t byte_offset, + int convolution_P, + int convolution_Q, + int add_P, + int add_Q, + int problem_N) const { + uint8_t* byte_pointer = byte_pointer_; + AccessType* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; + ++cluster) { + CUTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + CUTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + int frag_row_idx = + (row + + ThreadMap::Iterations::kRow * + (group + ThreadMap::Iterations::kGroup * cluster)); + + int row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + + bool row_guard = ((row_offset + thread_start_row_) < extent_row_); + + int output_row = row_offset + thread_start_row_; + int output_N = output_row / (convolution_P * convolution_Q); + int output_PQ = output_row % (convolution_P * convolution_Q); + int output_P = output_PQ / convolution_Q; + int output_Q = output_PQ % convolution_Q; + + int input_row = output_N * 2 * convolution_P * 2 * convolution_Q + + (2 * output_P + add_P) * 2 * convolution_Q + 2 * output_Q + add_Q; + + int64_t byte_offset = + (input_row - output_row) * problem_N * sizeof(float); + + AccessType* memory_pointer = + reinterpret_cast(byte_pointer + byte_offset); + + CUTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; + ++column) { + bool guard = row_guard && mask_.predicates[column]; + + cutlass::arch::global_load( + frag_ptr + [frag_row_idx * ThreadMap::Iterations::kColumn + column], + (void*)&memory_pointer + [column * ThreadMap::Delta::kColumn / kElementsPerAccess], + guard); + } + + if (row + 1 < ThreadMap::Iterations::kRow) { + byte_pointer += params_.increment_row; + } + } + + if (group + 1 < ThreadMap::Iterations::kGroup) { + byte_pointer += params_.increment_group; + } + } + + if (cluster + 1 < ThreadMap::Iterations::kCluster) { + byte_pointer += params_.increment_cluster; + } + } + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void upsample_load_with_byte_offset( + Fragment& frag, + int64_t byte_offset, + int convolution_P, + int convolution_Q, + int add_P, + int add_Q, + int problem_N) const { + uint8_t* byte_pointer = byte_pointer_; + AccessType* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; + ++cluster) { + CUTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + CUTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + int frag_row_idx = + (row + + ThreadMap::Iterations::kRow * + (group + ThreadMap::Iterations::kGroup * cluster)); + + int row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + + bool row_guard = ((row_offset + thread_start_row_) < extent_row_); + + int output_row = row_offset + thread_start_row_; + int output_N = output_row / (convolution_P * convolution_Q); + int output_PQ = output_row % (convolution_P * convolution_Q); + int output_P = output_PQ / convolution_Q; + int output_Q = output_PQ % convolution_Q; + int row_add_P = add_P; + int row_add_Q = add_Q; + if (output_P > convolution_P - 2) + row_add_P = 0; + if (output_Q > convolution_Q - 2) + row_add_Q = 0; + + int input_row = output_N * (convolution_P / 2) * (convolution_Q / 2) + + ((output_P + row_add_P) / 2) * (convolution_Q / 2) + + (output_Q + row_add_Q) / 2; + + int64_t byte_offset = + (input_row - output_row) * problem_N * sizeof(float); + + AccessType* memory_pointer = + reinterpret_cast(byte_pointer + byte_offset); + + CUTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; + ++column) { + bool guard = row_guard && mask_.predicates[column]; + + cutlass::arch::global_load( + frag_ptr + [frag_row_idx * ThreadMap::Iterations::kColumn + column], + (void*)&memory_pointer + [column * ThreadMap::Delta::kColumn / kElementsPerAccess], + guard); + } + + if (row + 1 < ThreadMap::Iterations::kRow) { + byte_pointer += params_.increment_row; + } + } + + if (group + 1 < ThreadMap::Iterations::kGroup) { + byte_pointer += params_.increment_group; + } + } + + if (cluster + 1 < ThreadMap::Iterations::kCluster) { + byte_pointer += params_.increment_cluster; + } + } + } + + CUTLASS_DEVICE + MatrixCoord thread_start() const { + return MatrixCoord(thread_start_row_, thread_start_column_); + } + + /// Need to get the thread start row from the tile iterator + CUTLASS_DEVICE + int32_t thread_start_row() const { + return thread_start_row_; + } + + /// Need to get the thread start row from the tile iterator + CUTLASS_DEVICE + int32_t thread_start_column() const { + return thread_start_column_; + } + + /// Extent of the matrix in rows + CUTLASS_DEVICE + Index extent_row() const { + return extent_row_; + } + + /// Extent of the matrix in columns + CUTLASS_DEVICE + Index extent_column() const { + return extent_column_; + } + + /// Advances to the next position to load or store + CUTLASS_HOST_DEVICE + PredicatedTileIteratorPrefetch& operator++() { + ++state_[0]; + + if (!ScatterD) { + byte_pointer_ += params_.advance_row; + } + + thread_start_row_ += ThreadMap::Shape::kRow; + + if (state_[0] == ThreadMap::Count::kRow) { + state_[0] = 0; + ++state_[1]; + byte_pointer_ += params_.advance_group; + + thread_start_row_ += (ThreadMap::Shape::kGroup - 1) * + ThreadMap::Shape::kRow * ThreadMap::Count::kRow; + + if (state_[1] == ThreadMap::Count::kGroup) { + state_[1] = 0; + ++state_[2]; + byte_pointer_ += params_.advance_cluster; + + thread_start_row_ += ThreadMap::Count::kGroup * + ThreadMap::Shape::kGroup * ThreadMap::Count::kRow * + ThreadMap::Shape::kRow; + + if (state_[2] == ThreadMap::Count::kCluster) { + state_[2] = 0; + byte_pointer_ += params_.advance_tile; + } + } + } + + return *this; + } + + ///< Efficiently disables all accesses guarded by mask + CUTLASS_DEVICE void clear_mask() { + mask_.clear(); + } + + ///< Efficiently enables all accesses guarded by mask + CUTLASS_DEVICE void enable_mask() { + mask_.enable(); + } + + ///< Sets the mask + CUTLASS_DEVICE void get_mask(Mask& mask) const { + mask = mask_; + } + + ///< Sets the mask + CUTLASS_DEVICE void set_mask(Mask const& mask) { + mask_ = mask; + } +}; + +template +struct MakePrefetchableIterator { + using Iterator = PredicatedTileIteratorPrefetch< + typename IT::ThreadMap, + typename IT::Element>; +}; + +/////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fmha/iterators/make_residual_last.h b/natten/csrc/include/natten/cuda/fmha/iterators/make_residual_last.h new file mode 100644 index 00000000..2930f063 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha/iterators/make_residual_last.h @@ -0,0 +1,73 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ +#pragma once + +#include +#include + +namespace cutlass { +namespace transform { +namespace threadblock { + +template +struct MakeIteratorResidualLast; + +template < + typename Shape, + typename Element, + typename Layout, + int AdvanceRank, + typename ThreadMap, + int AccessSize, + bool Gather> +struct MakeIteratorResidualLast> { + using Iterator = PredicatedTileIteratorResidualLast< + Shape, + Element, + Layout, + AdvanceRank, + ThreadMap, + AccessSize, + Gather>; +}; + +template < + typename Shape, + typename Element, + typename Layout, + int AdvanceRank, + typename ThreadMap, + typename AccessType, + bool Gather> +struct MakeIteratorResidualLast> { + using Iterator = PredicatedTileAccessIteratorResidualLast< + Shape, + Element, + Layout, + AdvanceRank, + ThreadMap, + AccessType, + Gather>; +}; +} // namespace threadblock +} // namespace transform +} // namespace cutlass diff --git a/natten/csrc/include/natten/cuda/fmha/iterators/predicated_tile_access_iterator_residual_last.h b/natten/csrc/include/natten/cuda/fmha/iterators/predicated_tile_access_iterator_residual_last.h new file mode 100644 index 00000000..ac08ea03 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha/iterators/predicated_tile_access_iterator_residual_last.h @@ -0,0 +1,2122 @@ +/* + * Copied from xFormers (https://github.com/facebookresearch/xformers/) and + * edited. + * + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + */ +/*************************************************************************************************** + * Copyright (c) 2017 - 2026 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 Templates calculating the address and predicates to the load of tiles + from pitch-linear rank=2 tensors. + + This iterator uses masks to guard out-of-bounds accesses. The first tile + this iterator visits maybe partial, then the remaining tiles are complete. + So, we only need to compute the predicates twice, once before the first tile + and once for the remaining full tiles which can share the same predicates. + + A precomputed "Params" object minimizes the amount of state that must be + stored in registers, and integer addition is used to advance the pointer + through memory. +*/ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +//////////////////////////////////////////////////////////////////////////////// + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace transform { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// PredicatedTileAccessIteratorResidualLast +/// +template < + typename Shape, + typename Element, + typename Layout, + int AdvanceRank, + typename ThreadMap, + typename AccessType, + bool Gather = false> +class PredicatedTileAccessIteratorResidualLast; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIteratorResidualLast for pitch-linear +/// data. +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + typename AccessType_, + bool Gather> +class PredicatedTileAccessIteratorResidualLast< + Shape_, + Element_, + layout::PitchLinear, + AdvanceRank, + ThreadMap_, + AccessType_, + Gather> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::PitchLinear; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingPredicates = PredicatedTileAccessIteratorPredicates< + Shape, + Element, + Layout, + AdvanceRank, + ThreadMap, + AccessType>; + + static int const kAccessesPerVector = + ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert( + !(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + using Mask = typename UnderlyingPredicates::Mask; + + /// Uses a non-template class + struct Params : PredicatedTileAccessIteratorParams { + using Base = PredicatedTileAccessIteratorParams; + + // Default ctor + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) + : Base( + layout.stride(0), + MakePredicatedTileAccessIteratorDesc< + Shape, + Element, + Layout, + kAdvanceRank, + ThreadMap>()()) {} + + CUTLASS_HOST_DEVICE + Params(Base const& base) : Base(base) {} + }; + + private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char*; + + private: + // + // Data members + // + + UnderlyingPredicates the_predicates; + Mask residual_tile_mask; + + /// Parameters object with precomputed internal state + Params params_; + + /// Internal pointer to first access of tile + BytePointer pointer_; + + /// Below is used when Gather is turned on. We need to record strided_offset + /// and contiguous_offset seperated to compute the offset by using + /// + /// offset = contiguous_offset + indices[strided_offset] + /// + + /// Gather indices + int const* indices_; + + Index gather_offset_strided; + + private: + /// Computes predicates based on internally tracked per-thread offset. + CUTLASS_DEVICE + void compute_predicates_( + /// Extent of the matrix window + TensorCoord extent, + /// optionally, simplify predicate calculation during 'steady state' phase + bool is_steady_state = false) { + the_predicates.compute_predicates_(extent, is_steady_state); + } + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const& threadblock_offset, + /// Gather indices + int const* indices = nullptr) + : params_(params), + pointer_(reinterpret_cast( + const_cast(pointer))), + the_predicates(extent), + indices_(indices) { + the_predicates.set_predicates(thread_id, threadblock_offset); + the_predicates.get_mask(residual_tile_mask); + + // Working around a weird compiler bug happening on P100 for the backward. + // I've seen together: the_predicates.predicates_[0] = 14 (instead of 15) + // residual_tile_mask[0] = 15 (correct) + // + // Adding prints when the value is calculated (in `compute_predicates_`) + // sometimes removes the bug. The consequence is that we skip some + // element of a tensor, leading to wrong results + // Setting `compute_predicates_`'s second argument (`is_steady_state`) to + // true also seems to get rid of the bug - at the cost of twice as many + // comparisons. +#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 700) + constexpr bool kWorkAroundCompilerBug = false; +#else + constexpr bool kWorkAroundCompilerBug = true; +#endif + the_predicates.compute_predicates_(extent, true && !kWorkAroundCompilerBug); + + // update internal pointers + Layout layout(params_.stride_); + + if (!Gather) { + add_pointer_offset(layout(the_predicates.thread_offset_)); + } else { + gather_offset_strided = the_predicates.thread_offset_.strided(); + add_pointer_offset( + layout(make_Coord(the_predicates.thread_offset_.contiguous(), 0))); + } + } + + /// Construct a PredicatedTileAccessIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + ///< ID of each participating thread + int thread_id) + : PredicatedTileAccessIteratorResidualLast( + params, + pointer, + extent, + thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) { + the_predicates.set_iteration_index(index); + } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool is_residual_tile) { + if (is_residual_tile) { + the_predicates.set_mask(residual_tile_mask); + } + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += sizeof_bits::value * pointer_offset / 8; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_DEVICE + void add_tile_offset(TensorCoord const& tile_offset) { + if (!Gather) { + if (kAdvanceRank) { + pointer_ += params_.inc_advance_ * LongIndex(tile_offset.strided()); + pointer_ += Shape::kContiguous * tile_offset.contiguous(); + } else { + pointer_ += params_.inc_advance_ * LongIndex(tile_offset.contiguous()); + pointer_ += Shape::kStrided * tile_offset.strided(); + } + } else { + add_pointer_offset(Shape::kContiguous * tile_offset.contiguous()); + gather_offset_strided += Shape::kStrided * tile_offset.strided(); + } + } + + /// Returns a pointer + CUTLASS_HOST_DEVICE + AccessType* get() const { + if (Gather) { + assert(indices_); + + if (!valid()) { + return nullptr; + } + + LongIndex contiguous_offset = the_predicates.iteration_contiguous_ * + (ThreadMap::Delta::kContiguous * sizeof_bits::value / + 8) + + the_predicates.iteration_vector_; + int strided_index = gather_offset_strided + + the_predicates.iteration_strided_ * ThreadMap::Delta::kStrided; + + LongIndex strided_offset = indices_[strided_index] * + LongIndex(params_.stride_) * sizeof_bits::value / 8; + + return reinterpret_cast( + pointer_ + contiguous_offset + strided_offset); + } + + return reinterpret_cast( + pointer_ + + the_predicates.iteration_contiguous_ * + (ThreadMap::Delta::kContiguous * + sizeof_bits::value) / + 8) + + the_predicates.iteration_vector_; + } + + /// Increment and return an instance to self. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast& operator++() { + the_predicates.operator++(); + + ++the_predicates.iteration_vector_; + if (the_predicates.iteration_vector_ < kAccessesPerVector) { + return *this; + } + + the_predicates.iteration_vector_ = 0; + ++the_predicates.iteration_contiguous_; + + if (the_predicates.iteration_contiguous_ < + ThreadMap::Iterations::kContiguous) { + return *this; + } + + // Enter here only if (iteration_contiguous_ == + // ThreadMap::Iteration::kContiguous) + the_predicates.iteration_contiguous_ = 0; + ++the_predicates.iteration_strided_; + + if (the_predicates.iteration_strided_ < ThreadMap::Iterations::kStrided) { + if (!Gather) { + pointer_ += params_.inc_strided_; + } + + return *this; + } + + // Enter here only if (iteration_stride_ == ThreadMap::Iteration::kStrided) + // which means we enter the next tile. + the_predicates.iteration_strided_ = 0; + + if (!Gather) { + // advance to next tile + pointer_ += params_.inc_next_; + + // now return to start tile - if the iterator is subsequently advanced, + // this subtraction as well as the subsequent integer addition are both + // elided by the compiler. + pointer_ -= params_.inc_advance_; + } + + return *this; + } + + /// Increment and return an instance to self. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast operator++(int) { + PredicatedTileAccessIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + the_predicates.clear_mask(enable); + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { + the_predicates.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { + the_predicates.set_mask(mask); + } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { + the_predicates.get_mask(mask); + } + + /// Returns whether access is valid or not + CUTLASS_HOST_DEVICE + bool valid() const { + return the_predicates.valid(); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIteratorResidualLast for column-major +/// data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + typename AccessType_, + bool Gather> +class PredicatedTileAccessIteratorResidualLast< + Shape_, + Element_, + layout::ColumnMajor, + AdvanceRank, + ThreadMap_, + AccessType_, + Gather> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingIterator = PredicatedTileAccessIteratorResidualLast< + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 0 : 1), + ThreadMap, + AccessType, + Gather>; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileAccessIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + /// Default ctor + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) + : params_(layout::PitchLinear(layout.stride(0))){}; + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const& base) + : params_(base) {} + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + ///< Precomputed parameters object + Params const& params, + ///< Pointer to start of tensor + Pointer pointer, + ///< Extent of tensor + TensorCoord extent, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const& threadblock_offset, + int const* indices = + nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : iterator_( + params.params_, + pointer, + layout::PitchLinearCoord(extent.row(), extent.column()), + thread_id, + layout::PitchLinearCoord( + threadblock_offset.row(), + threadblock_offset.column()), + indices) {} + + /// Construct a PredicatedTileAccessIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileAccessIteratorResidualLast( + params, + pointer, + extent, + thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) { + iterator_.set_iteration_index(index); + } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { + iterator_.set_residual_tile(enable); + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const& tile_offset) { + iterator_.add_tile_offset({tile_offset.row(), tile_offset.column()}); + } + + /// Returns a pointer + CUTLASS_HOST_DEVICE + AccessType* get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast& operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast operator++(int) { + PredicatedTileAccessIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + iterator_.clear_mask(enable); + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { + iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { + iterator_.set_mask(mask); + } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { + iterator_.get_mask(mask); + } + + /// Returns whether access is valid or not + CUTLASS_HOST_DEVICE + bool valid() { + return iterator_.valid(); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIteratorResidualLast for row-major +/// data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + typename AccessType_, + bool Gather> +class PredicatedTileAccessIteratorResidualLast< + Shape_, + Element_, + layout::RowMajor, + AdvanceRank, + ThreadMap_, + AccessType_, + Gather> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingIterator = PredicatedTileAccessIteratorResidualLast< + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap, + AccessType, + Gather>; + + static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileAccessIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + /// Default ctor + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) + : params_(layout::PitchLinear(layout.stride(0))){}; + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const& base) + : params_(base) {} + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + ///< Precomputed parameters object + Params const& params, + ///< Pointer to start of tensor + Pointer pointer, + ///< Extent of tensor + TensorCoord extent, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const& threadblock_offset, + /// Gather indices + int const* indices = nullptr) + : iterator_( + params.params_, + pointer, + layout::PitchLinearCoord(extent.column(), extent.row()), + thread_id, + layout::PitchLinearCoord( + threadblock_offset.column(), + threadblock_offset.row()), + indices) {} + + /// Construct a PredicatedTileAccessIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileAccessIteratorResidualLast( + params, + pointer, + extent, + thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) { + iterator_.set_iteration_index(index); + } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { + iterator_.set_residual_tile(enable); + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const& tile_offset) { + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + } + + /// Returns a pointer + CUTLASS_HOST_DEVICE + AccessType* get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast& operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast operator++(int) { + PredicatedTileAccessIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + iterator_.clear_mask(enable); + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { + iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { + iterator_.set_mask(mask); + } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { + iterator_.get_mask(mask); + } + + /// Returns whether access is valid or not + CUTLASS_HOST_DEVICE + bool valid() { + return iterator_.valid(); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIteratorResidualLast for affine rank 2 +/// data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + typename AccessType_> +class PredicatedTileAccessIteratorResidualLast< + Shape_, + Element_, + layout::AffineRankN<2>, + AdvanceRank, + ThreadMap_, + AccessType_, + false> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::AffineRankN<2>; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingPredicates = PredicatedTileAccessIteratorPredicates< + Shape, + Element, + layout::PitchLinear, + AdvanceRank, + ThreadMap, + AccessType>; + + static int const kAccessesPerVector = + ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert( + !(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingPredicates::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + public: + friend PredicatedTileAccessIteratorResidualLast; + + private: + /// stride of pitch-linear layout (units of Element) + Coord stride_; + /// amount (in byte) to increment pointer to move to next access along + /// contiguous dimension + LongIndex inc_contiguous_; + /// amount (in byte) to increment pointer from first access of current + /// contiguous dimension to first access of next one. + LongIndex inc_strided_; + /// amount (in byte) to increment pointer from last access of current + /// contiguous dimension to first access of next one. + LongIndex inc_next_strided_; + /// amount (in byte) to increment pointer from last access to first access + /// of next tile + LongIndex inc_next_; + /// amount (in byte) to increment pointer from first access of current tile + /// to first access of next tile + LongIndex inc_advance_; + + public: + // Default ctor + CUTLASS_HOST_DEVICE + Params() + : stride_(0), + inc_contiguous_(0), + inc_strided_(0), + inc_next_(0), + inc_advance_(0) {} + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) + : stride_({layout.stride(0), layout.stride(1)}) { + inc_contiguous_ = + (LongIndex(stride_[0]) * ThreadMap::Delta::kContiguous) * + sizeof_bits::value / 8; + + inc_strided_ = (LongIndex(stride_[1]) * ThreadMap::Delta::kStrided) * + sizeof_bits::value / 8; + + inc_next_strided_ = inc_strided_ - + LongIndex(ThreadMap::Iterations::kContiguous - 1) * inc_contiguous_; + + if (kAdvanceRank) { + // advance along strided dimension + inc_advance_ = Shape::kStrided * LongIndex(stride_[1]) * + sizeof_bits::value / 8; + } else { + // advance along contiguous dimension + inc_advance_ = + Shape::kContiguous * stride_[0] * sizeof_bits::value / 8; + } + + inc_next_ = inc_advance_ - + LongIndex(ThreadMap::Iterations::kContiguous - 1) * inc_contiguous_ - + LongIndex(ThreadMap::Iterations::kStrided - 1) * inc_strided_; + }; + }; + + private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char*; + + // + // Data members + // + + /// Parameters object with precomputed internal state + Params params_; + + /// Internal pointer to first access of tile + BytePointer pointer_; + + UnderlyingPredicates the_predicates; + Mask residual_tile_mask; + + private: + /// Computes predicates based on internally tracked per-thread offset. + CUTLASS_DEVICE + void compute_predicates_( + /// Extent of the matrix window + TensorCoord extent, + /// optionally, simplify predicate calculation during 'steady state' phase + bool is_steady_state = false) { + the_predicates.compute_predicates_(extent, is_steady_state); + } + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + ///< Precomputed parameters object + Params const& params, + ///< Pointer to start of tensor + Pointer pointer, + ///< Extent of tensor + TensorCoord extent, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const& threadblock_offset, + int const* indices = + nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : params_(params), + pointer_(reinterpret_cast( + const_cast(pointer))), + the_predicates(extent) { + the_predicates.set_predicates(thread_id, threadblock_offset); + + // update internal pointers + Layout layout(params_.stride_); + add_pointer_offset(layout(the_predicates.thread_offset_)); + } + + /// Construct a PredicatedTileAccessIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileAccessIteratorResidualLast( + params, + pointer, + extent, + thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) { + the_predicates.set_iteration_index(index); + } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool is_residual_tile) { + if (is_residual_tile) { + the_predicates.set_mask(residual_tile_mask); + } + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += sizeof_bits::value * pointer_offset / 8; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const& tile_offset) { + if (kAdvanceRank) { + pointer_ += params_.inc_advance_ * LongIndex(tile_offset[1]); + pointer_ += Shape::kContiguous * tile_offset[0]; + } else { + pointer_ += params_.inc_advance_ * LongIndex(tile_offset[0]); + pointer_ += Shape::kStrided * tile_offset[1]; + } + } + + /// Returns a pointer + CUTLASS_HOST_DEVICE + AccessType* get() const { + return reinterpret_cast(pointer_) + + the_predicates.iteration_vector_; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast& operator++() { + the_predicates.operator++(); + ++the_predicates.iteration_vector_; + if (the_predicates.iteration_vector_ < kAccessesPerVector) { + return *this; + } + + the_predicates.iteration_vector_ = 0; + ++the_predicates.iteration_contiguous_; + + if (the_predicates.iteration_contiguous_ < + ThreadMap::Iterations::kContiguous) { + pointer_ += params_.inc_contiguous_; + return *this; + } + + // Enter here only if (iteration_contiguous_ == + // ThreadMap::Iteration::kContiguous) + the_predicates.iteration_contiguous_ = 0; + ++the_predicates.iteration_strided_; + + if (the_predicates.iteration_strided_ < ThreadMap::Iterations::kStrided) { + pointer_ += params_.inc_next_strided_; + return *this; + } + + // Enter here only if (iteration_stride_ == ThreadMap::Iteration::kStrided) + // which means we enter the next tile. + the_predicates.iteration_strided_ = 0; + + // advance to next tile + pointer_ += params_.inc_next_; + + // now return to start tile - if the iterator is subsequently advanced, this + // subtraction as well as the subsequent integer addition are both elided by + // the compiler. + pointer_ -= params_.inc_advance_; + + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast operator++(int) { + PredicatedTileAccessIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + the_predicates.clear_mask(enable); + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { + the_predicates.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { + the_predicates.set_mask(mask); + } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { + the_predicates.get_mask(mask); + } + + /// Returns whether access is valid or not + CUTLASS_HOST_DEVICE + bool valid() { + return the_predicates.valid(); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIteratorResidualLast for affine rank 2 +/// column-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + typename AccessType_> +class PredicatedTileAccessIteratorResidualLast< + Shape_, + Element_, + layout::AffineRank2ColumnMajor, + AdvanceRank, + ThreadMap_, + AccessType_, + false> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::AffineRank2ColumnMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + // Map to the underlying AffineRankN<2> layout + using UnderlyingIterator = PredicatedTileAccessIteratorResidualLast< + layout::PitchLinearShape, + Element, + layout::AffineRankN<2>, + (kAdvanceRank == 0 ? 0 : 1), + ThreadMap, + AccessType>; + + static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileAccessIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + /// Default ctor + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given an AffineRankN<2> tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) + : params_(layout::AffineRankN<2>(layout.stride(0), layout.stride(1))){}; + }; + + private: + // + // Data members + // + + /// Underlying AffineRankN<2> tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + ///< Precomputed parameters object + Params const& params, + ///< Pointer to start of tensor + Pointer pointer, + ///< Extent of tensor + TensorCoord extent, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const& threadblock_offset, + int const* indices = + nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : iterator_( + params.params_, + pointer, + layout::PitchLinearCoord(extent.row(), extent.column()), + thread_id, + layout::PitchLinearCoord( + threadblock_offset.row(), + threadblock_offset.column())) {} + + /// Construct a PredicatedTileAccessIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileAccessIteratorResidualLast( + params, + pointer, + extent, + thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) { + iterator_.set_iteration_index(index); + } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { + iterator_.set_residual_tile(enable); + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const& tile_offset) { + iterator_.add_tile_offset( + make_Coord(tile_offset.row(), tile_offset.column())); + } + + /// Returns a pointer + CUTLASS_HOST_DEVICE + AccessType* get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast& operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast operator++(int) { + PredicatedTileAccessIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + iterator_.clear_mask(enable); + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { + iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { + iterator_.set_mask(mask); + } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { + iterator_.get_mask(mask); + } + + /// Returns whether access is valid or not + CUTLASS_HOST_DEVICE + bool valid() { + return iterator_.valid(); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIteratorResidualLast for affine rank-2 +/// row-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + typename AccessType_> +class PredicatedTileAccessIteratorResidualLast< + Shape_, + Element_, + layout::AffineRank2RowMajor, + AdvanceRank, + ThreadMap_, + AccessType_, + false> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::AffineRank2RowMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + // Map to the underlying AffineRankN<2> layout + using UnderlyingIterator = PredicatedTileAccessIteratorResidualLast< + layout::PitchLinearShape, + Element, + layout::AffineRankN<2>, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap, + AccessType>; + + static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileAccessIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + /// Default ctor + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given an AffineRankN<2> tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) + : params_(layout::AffineRankN<2>(layout.stride(1), layout.stride(0))){}; + }; + + private: + // + // Data members + // + + /// Underlying AffineRankN<2> tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + ///< Precomputed parameters object + Params const& params, + ///< Pointer to start of tensor + Pointer pointer, + ///< Extent of tensor + TensorCoord extent, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const& threadblock_offset, + int const* indices = + nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : iterator_( + params.params_, + pointer, + layout::PitchLinearCoord(extent.column(), extent.row()), + thread_id, + layout::PitchLinearCoord( + threadblock_offset.column(), + threadblock_offset.row())) {} + + /// Construct a PredicatedTileAccessIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileAccessIteratorResidualLast( + params, + pointer, + extent, + thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) { + iterator_.set_iteration_index(index); + } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { + iterator_.set_residual_tile(enable); + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const& tile_offset) { + iterator_.add_tile_offset( + make_Coord(tile_offset.column(), tile_offset.row())); + } + + /// Returns a pointer + CUTLASS_HOST_DEVICE + AccessType* get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast& operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast operator++(int) { + PredicatedTileAccessIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + iterator_.clear_mask(enable); + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { + iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { + iterator_.set_mask(mask); + } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { + iterator_.get_mask(mask); + } + + /// Returns whether access is valid or not + CUTLASS_HOST_DEVICE + bool valid() { + return iterator_.valid(); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIteratorResidualLast for column-major +/// interleaved data. It is mapped to the congruous layout. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// + +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + typename AccessType_, + int InterleavedK> +class PredicatedTileAccessIteratorResidualLast< + Shape_, + Element_, + layout::ColumnMajorInterleaved, + AdvanceRank, + ThreadMap_, + AccessType_, + false> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + static int const kInterleavedK = InterleavedK; + using Layout = layout::ColumnMajorInterleaved; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingIterator = PredicatedTileAccessIteratorResidualLast< + layout::PitchLinearShape< + Shape::kRow * kInterleavedK, + Shape::kColumn / kInterleavedK>, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 0 : 1), + ThreadMap, + AccessType>; + + static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileAccessIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) + : params_(layout::PitchLinear(layout.stride(0))) {} + + CUTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const& base) + : params_(base) {} + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const& threadblock_offset, + int const* indices = + nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : iterator_( + params.params_, + pointer, + layout::PitchLinearCoord( + extent.row() * kInterleavedK, + extent.column() / kInterleavedK), + thread_id, + layout::PitchLinearCoord( + threadblock_offset.row() * kInterleavedK, + threadblock_offset.column() / kInterleavedK)) {} + + /// Construct a PredicatedTileAccessIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileAccessIteratorResidualLast( + params, + pointer, + extent, + thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) { + iterator_.set_iteration_index(index); + } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { + iterator_.set_residual_tile(enable); + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const& tile_offset) { + iterator_.add_tile_offset({tile_offset.row(), tile_offset.column()}); + } + + /// Returns a pointer + CUTLASS_HOST_DEVICE + AccessType* get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast& operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast operator++(int) { + PredicatedTileAccessIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + iterator_.clear_mask(enable); + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { + iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { + iterator_.set_mask(mask); + } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { + iterator_.get_mask(mask); + } + + /// Returns whether access is valid or not + CUTLASS_HOST_DEVICE + bool valid() { + return iterator_.valid(); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIteratorResidualLast for row-major +/// interleaved data. +// It is mapped to the congruous layout. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + typename AccessType_, + int InterleavedK> +class PredicatedTileAccessIteratorResidualLast< + Shape_, + Element_, + layout::RowMajorInterleaved, + AdvanceRank, + ThreadMap_, + AccessType_, + false> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + static int const kInterleavedK = InterleavedK; + using Layout = layout::RowMajorInterleaved; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingIterator = PredicatedTileAccessIteratorResidualLast< + layout::PitchLinearShape< + Shape::kColumn * kInterleavedK, + Shape::kRow / kInterleavedK>, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap, + AccessType>; + + static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileAccessIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) + : params_(layout::PitchLinear(layout.stride(0))) {} + + CUTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const& base) + : params_(base) {} + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const& threadblock_offset, + int const* indices = + nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : iterator_( + params.params_, + pointer, + layout::PitchLinearCoord( + extent.column() * kInterleavedK, + extent.row() / kInterleavedK), + thread_id, + layout::PitchLinearCoord( + threadblock_offset.column() * kInterleavedK, + threadblock_offset.row() / kInterleavedK)) {} + + /// Construct a PredicatedTileAccessIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileAccessIteratorResidualLast( + params, + pointer, + extent, + thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) { + iterator_.set_iteration_index(index); + } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { + iterator_.set_residual_tile(enable); + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const& tile_offset) { + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + } + + /// Returns a pointer + CUTLASS_HOST_DEVICE + AccessType* get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast& operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast operator++(int) { + PredicatedTileAccessIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + iterator_.clear_mask(enable); + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { + iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { + iterator_.set_mask(mask); + } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { + iterator_.get_mask(mask); + } + + /// Returns whether access is valid or not + CUTLASS_HOST_DEVICE + bool valid() { + return iterator_.valid(); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fmha/iterators/predicated_tile_iterator_residual_last.h b/natten/csrc/include/natten/cuda/fmha/iterators/predicated_tile_iterator_residual_last.h new file mode 100644 index 00000000..ae4052cd --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha/iterators/predicated_tile_iterator_residual_last.h @@ -0,0 +1,2127 @@ +/* + * Copied from xFormers (https://github.com/facebookresearch/xformers/) and + * edited. + * + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + */ +/*************************************************************************************************** + * Copyright (c) 2017 - 2026 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 Templates implementing loading of tiles from pitch-linear rank=2 + tensors. + + This iterator uses masks to guard out-of-bounds accesses. The first tile + this iterator visits maybe partial, then the remaining tiles are complete. + So, we only need to compute the predicates twice, once before the first tile + and once for the remaining full tiles which can share the same predicates. + + A precomputed "Params" object minimizes the amount of state that must be + stored in registers, and integer addition is used to advance the pointer + through memory. +*/ + +#pragma once + +#include +#include + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace transform { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// PredicatedTileIteratorResidualLast +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +/// Regular tile iterator using a precomputed control structure to minimize +/// register liveness and integer arithmetic. +/// +/// Layout is assumed to be invariant at the time the precomputed "Params" +/// object is constructed. +/// +/// Base pointer and tensor extents may be specified at the time the iterator is +/// constructed. Subsequently, they are assumed to be immutable. +/// +/// Adding a logical coordinate offset may be performed at the time the iterator +/// is constructed. Subsequent additions to logical coordinate offset may be +/// performed but are relatively expensive. +/// +/// Visitation order is intended to first visit a "residual" tile that may be +/// partially full in both the advance dimension and the steady-state dimension. +/// This is assumed to be the last tile in the iteration sequence. Advancing an +/// iterator that has just been constructed moves to the first tile that is full +/// in the advance dimension and recomputes predicates. Subsequent accesses may +/// be performed without updating internal predicates and are efficient in terms +/// of live register state and pointer arithmetic instructions. +/// +/// To be efficient, this assumes the iterator will be dereferenced and advanced +/// at least once outside any looping structure to minimize integer arithmetic. +/// +/// Acceses out of bounds are safe so long as `clear_mask()` is called prior to +/// dereferencing the iterator. +/// +/// +/// Example: +/// +/// An efficient pipeline structure may be constructed as follows: +/// +// template +// __global__ void kernel( +// typename Iterator::Params params, +// typename Iterator::Element *ptr, +// TensorCoord extent) { +// +// typename Iterator::Fragment fragment; +// +// TensorCoord threadblock_offset(0, 0); +// +// Iterator iter(params, ptr, extent, threadIdx.x, threadblock_offsets); +// +// +// fragment = *iter; // load "residue" tile first +// ++iter; // advance to first "steady state" tile and update +// internal masks +// +// +// #pragma unroll +// for (int i = Remaining - 1; i >= 0; --i) { +// +// f(fragment); +// +// if (!i) { +// iter.clear_mask(); // light-weight operation to clear masks - +// subsequent loads become NO-OPs. +// } +// +// fragment = *iter; // load tile during "steady state" phase +// ++iter; // advance to next tile - lightweight due to +// steady-state masks +// } +// } +// +// void host(TensorView view) { +// +// using Iterator = +// transform::threadblock::PredicatedTileIteratorResidualLast; +// +// typename Iterator::Params params(view.layout()); +// +// kernel(params, view.data()); +// } +/// +/// +template < + typename Shape, + typename Element, + typename Layout, + int AdvanceRank, + typename ThreadMap, + int AccessSize = ThreadMap::kElementsPerAccess, + bool Gather = false> +class PredicatedTileIteratorResidualLast; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIteratorResidualLast for pitch-linear data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int AccessSize, + bool Gather> +class PredicatedTileIteratorResidualLast< + Shape_, + Element_, + layout::PitchLinear, + AdvanceRank, + ThreadMap_, + AccessSize, + Gather> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::PitchLinear; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + /// Type used for internal memory accesses + using AccessType = AlignedArray< + Element, + AccessSize, + (AccessSize * sizeof_bits::value / 8)>; + + /// Underlying iterator to compute the addresses + using TileAccessIterator = PredicatedTileAccessIteratorResidualLast< + Shape, + Element, + Layout, + kAdvanceRank, + ThreadMap, + AccessType, + Gather>; + + static int const kAccessesPerVector = TileAccessIterator::kAccessesPerVector; + + /// Fragment object to be loaded or stored + using Fragment = cutlass::Array< + Element, + ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; + + /// Predicate vector stores mask to guard accesses + using Mask = typename TileAccessIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + public: + using Base = typename TileAccessIterator::Params::Base; + + friend PredicatedTileIteratorResidualLast; + + private: + /// Parameters object + typename TileAccessIterator::Params params_; + + public: + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) : params_(layout) {} + + CUTLASS_HOST_DEVICE + Params() {} + + CUTLASS_HOST_DEVICE + Params(Base const& base) : params_(base) {} + }; + + private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char*; + + private: + // + // Data members + // + + /// Data member to the tile access iterator + TileAccessIterator address_iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const& threadblock_offset, + /// Gather indices + int const* indices = nullptr) + : address_iterator_( + params.params_, + pointer, + extent, + thread_id, + threadblock_offset, + indices) {} + + /// Construct a PredicatedTileIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIteratorResidualLast( + params, + pointer, + extent, + thread_id, + make_Coord(0, 0)) {} + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + address_iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast& operator++() { + if (kAdvanceRank) + address_iterator_.add_tile_offset({0, 1}); + else + address_iterator_.add_tile_offset({1, 0}); + + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast operator++(int) { + PredicatedTileIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + address_iterator_.clear_mask(enable); + } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { + address_iterator_.set_residual_tile(enable); + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { + address_iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { + address_iterator_.set_mask(mask); + } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { + address_iterator_.get_mask(mask); + } + + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment& frag, Index pointer_offset) { + load_with_byte_offset( + frag, pointer_offset * sizeof_bits::value / 8); + } + + CUTLASS_DEVICE + void load_with_byte_offset(Fragment& frag, LongIndex byte_offset) { + AccessType* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + int idx = v + + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); + + address_iterator_.set_iteration_index(idx); + char const* byte_ptr = + reinterpret_cast(address_iterator_.get()) + + byte_offset; + + AccessType const* access_ptr = + reinterpret_cast(byte_ptr); + + cutlass::arch::global_load( + frag_ptr[idx], access_ptr, address_iterator_.valid()); + + ++address_iterator_; + } + } + } + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) { + load_with_byte_offset(frag, 0); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) { + store_with_byte_offset( + frag, pointer_offset * sizeof_bits::value / 8); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, LongIndex byte_offset) { + address_iterator_.set_iteration_index(0); + AccessType const* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + int idx = v + + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); + + char* byte_ptr = + reinterpret_cast(address_iterator_.get()) + byte_offset; + AccessType* access_ptr = reinterpret_cast(byte_ptr); + + if (address_iterator_.valid()) { + *access_ptr = frag_ptr[idx]; + } + ++address_iterator_; + } + } + } + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { + store_with_byte_offset(frag, 0); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIteratorResidualLast for pitch-linear data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int AccessSize, + bool Gather> +class PredicatedTileIteratorResidualLast< + Shape_, + Element_, + layout::ColumnMajor, + AdvanceRank, + ThreadMap_, + AccessSize, + Gather> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingIterator = PredicatedTileIteratorResidualLast< + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 0 : 1), + ThreadMap, + AccessSize, + Gather>; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = cutlass::Array< + Element, + ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) + : params_(layout::PitchLinear(layout.stride(0))) {} + + CUTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const& base) + : params_(base) {} + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id, ///< ID of each participating thread + TensorCoord const& threadblock_offset, ///< Initial offset of threadblock + int const* indices = + nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : iterator_( + params.params_, + pointer, + layout::PitchLinearCoord(extent.row(), extent.column()), + thread_id, + layout::PitchLinearCoord( + threadblock_offset.row(), + threadblock_offset.column()), + indices) {} + + /// Construct a PredicatedTileIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIteratorResidualLast( + params, + pointer, + extent, + thread_id, + make_Coord(0, 0)) {} + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast& operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast operator++(int) { + PredicatedTileIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + iterator_.clear_mask(enable); + } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { + iterator_.set_residual_tile(enable); + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { + iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { + iterator_.set_mask(mask); + } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { + iterator_.get_mask(mask); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment& frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_byte_offset(Fragment& frag, LongIndex byte_offset) { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) { + load_with_pointer_offset(frag, 0); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, LongIndex byte_offset) { + iterator_.store_with_byte_offset(frag, byte_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { + store_with_pointer_offset(frag, 0); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIteratorResidualLast for pitch-linear data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int AccessSize, + bool Gather> +class PredicatedTileIteratorResidualLast< + Shape_, + Element_, + layout::RowMajor, + AdvanceRank, + ThreadMap_, + AccessSize, + Gather> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingIterator = PredicatedTileIteratorResidualLast< + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap, + AccessSize, + Gather>; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = cutlass::Array< + Element, + ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) + : params_(layout::PitchLinear(layout.stride(0))) {} + + CUTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const& base) + : params_(base) {} + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id, ///< ID of each participating thread + TensorCoord const& threadblock_offset, ///< Initial offset of threadblock + int const* indices = nullptr ///< Gather indices + ) + : iterator_( + params.params_, + pointer, + layout::PitchLinearCoord(extent.column(), extent.row()), + thread_id, + layout::PitchLinearCoord( + threadblock_offset.column(), + threadblock_offset.row()), + indices) {} + + /// Construct a PredicatedTileIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIteratorResidualLast( + params, + pointer, + extent, + thread_id, + make_Coord(0, 0)) {} + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast& operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast operator++(int) { + PredicatedTileIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + iterator_.clear_mask(enable); + } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { + iterator_.set_residual_tile(enable); + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { + iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { + iterator_.set_mask(mask); + } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { + iterator_.get_mask(mask); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment& frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_byte_offset(Fragment& frag, LongIndex byte_offset) { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) { + load_with_pointer_offset(frag, 0); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, LongIndex byte_offset) { + iterator_.store_with_byte_offset(frag, byte_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { + store_with_pointer_offset(frag, 0); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIteratorResidualLast for affine rank-2 data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int AccessSize> +class PredicatedTileIteratorResidualLast< + Shape_, + Element_, + layout::AffineRankN<2>, + AdvanceRank, + ThreadMap_, + AccessSize, + false> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::AffineRankN<2>; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + /// Type used for internal memory accesses + using AccessType = AlignedArray< + Element, + AccessSize, + (AccessSize * sizeof_bits::value / 8)>; + + /// Underlying iterator to compute the addresses + using TileAccessIterator = PredicatedTileAccessIteratorResidualLast< + Shape, + Element, + Layout, + kAdvanceRank, + ThreadMap, + AccessType>; + + static int const kAccessesPerVector = TileAccessIterator::kAccessesPerVector; + + /// Fragment object to be loaded or stored + using Fragment = cutlass::Array< + Element, + ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; + + /// Predicate vector stores mask to guard accesses + using Mask = typename TileAccessIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + public: + friend PredicatedTileIteratorResidualLast; + + private: + /// Parameters object + typename TileAccessIterator::Params params_; + + public: + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) : params_(layout) {} + + CUTLASS_HOST_DEVICE + Params() {} + }; + + private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char*; + + private: + // + // Data members + // + + /// Data member to the tile access iterator + TileAccessIterator address_iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const& threadblock_offset, + int const* indices = + nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : address_iterator_( + params.params_, + pointer, + extent, + thread_id, + threadblock_offset) {} + + /// Construct a PredicatedTileIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIteratorResidualLast( + params, + pointer, + extent, + thread_id, + make_Coord(0, 0)) {} + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + address_iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast& operator++() { + if (kAdvanceRank) + address_iterator_.add_tile_offset(make_Coord(0, 1)); + else + address_iterator_.add_tile_offset(make_Coord(1, 0)); + + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast operator++(int) { + PredicatedTileIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + address_iterator_.clear_mask(enable); + } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { + address_iterator_.set_residual_tile(enable); + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { + address_iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { + address_iterator_.set_mask(mask); + } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { + address_iterator_.get_mask(mask); + } + + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment& frag, Index pointer_offset) { + load_with_byte_offset( + frag, pointer_offset * sizeof_bits::value / 8); + } + + CUTLASS_DEVICE + void load_with_byte_offset(Fragment& frag, LongIndex byte_offset) { + AccessType* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + int idx = v + + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); + + address_iterator_.set_iteration_index(idx); + char const* byte_ptr = + reinterpret_cast(address_iterator_.get()) + + byte_offset; + + AccessType const* access_ptr = + reinterpret_cast(byte_ptr); + + cutlass::arch::global_load( + frag_ptr[idx], access_ptr, address_iterator_.valid()); + + ++address_iterator_; + } + } + } + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) { + load_with_byte_offset(frag, 0); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) { + store_with_byte_offset( + frag, pointer_offset * sizeof_bits::value / 8); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, LongIndex byte_offset) { + address_iterator_.set_iteration_index(0); + AccessType const* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + int idx = v + + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); + + char* byte_ptr = + reinterpret_cast(address_iterator_.get()) + byte_offset; + AccessType* access_ptr = reinterpret_cast(byte_ptr); + + if (address_iterator_.valid()) { + *access_ptr = frag_ptr[idx]; + } + ++address_iterator_; + } + } + } + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { + store_with_byte_offset(frag, 0); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIteratorResidualLast for affine rank 2 +/// column-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int AccessSize> +class PredicatedTileIteratorResidualLast< + Shape_, + Element_, + layout::AffineRank2ColumnMajor, + AdvanceRank, + ThreadMap_, + AccessSize, + false> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::AffineRank2ColumnMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + // Map to the underlying AffineRankN<2> layout + using UnderlyingIterator = PredicatedTileIteratorResidualLast< + layout::PitchLinearShape, + Element, + layout::AffineRankN<2>, + (kAdvanceRank == 0 ? 0 : 1), + ThreadMap, + AccessSize>; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = cutlass::Array< + Element, + ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given an AffineRankN<2> tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) + : params_(layout::AffineRankN<2>(layout.stride(0), layout.stride(1))) {} + }; + + private: + // + // Data members + // + + /// Underlying AffineRankN<2> tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id, ///< ID of each participating thread + TensorCoord const& threadblock_offset, ///< Initial offset of threadblock + int const* indices = + nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : iterator_( + params.params_, + pointer, + layout::PitchLinearCoord(extent.row(), extent.column()), + thread_id, + layout::PitchLinearCoord( + threadblock_offset.row(), + threadblock_offset.column())) {} + + /// Construct a PredicatedTileIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIteratorResidualLast( + params, + pointer, + extent, + thread_id, + make_Coord(0, 0)) {} + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast& operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast operator++(int) { + PredicatedTileIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + iterator_.clear_mask(enable); + } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { + iterator_.set_residual_tile(enable); + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { + iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { + iterator_.set_mask(mask); + } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { + iterator_.get_mask(mask); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment& frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_byte_offset(Fragment& frag, LongIndex byte_offset) { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) { + load_with_pointer_offset(frag, 0); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, LongIndex byte_offset) { + iterator_.store_with_byte_offset(frag, byte_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { + store_with_pointer_offset(frag, 0); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIteratorResidualLast for affine rank 2 +/// row-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int AccessSize> +class PredicatedTileIteratorResidualLast< + Shape_, + Element_, + layout::AffineRank2RowMajor, + AdvanceRank, + ThreadMap_, + AccessSize, + false> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::AffineRank2RowMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + // Map to the underlying AffineRankN<2> layout + using UnderlyingIterator = PredicatedTileIteratorResidualLast< + layout::PitchLinearShape, + Element, + layout::AffineRankN<2>, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap, + AccessSize>; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = cutlass::Array< + Element, + ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given an AffineRankN<2> tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) + : params_(layout::AffineRankN<2>(layout.stride(1), layout.stride(0))) {} + }; + + private: + // + // Data members + // + + /// Underlying AffineRankN<2> tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id, ///< ID of each participating thread + TensorCoord const& threadblock_offset, ///< Initial offset of threadblock + int const* indices = + nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : iterator_( + params.params_, + pointer, + layout::PitchLinearCoord(extent.column(), extent.row()), + thread_id, + layout::PitchLinearCoord( + threadblock_offset.column(), + threadblock_offset.row())) {} + + /// Construct a PredicatedTileIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIteratorResidualLast( + params, + pointer, + extent, + thread_id, + make_Coord(0, 0)) {} + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast& operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast operator++(int) { + PredicatedTileIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + iterator_.clear_mask(enable); + } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { + iterator_.set_residual_tile(enable); + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { + iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { + iterator_.set_mask(mask); + } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { + iterator_.get_mask(mask); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment& frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_byte_offset(Fragment& frag, LongIndex byte_offset) { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) { + load_with_pointer_offset(frag, 0); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, LongIndex byte_offset) { + iterator_.store_with_byte_offset(frag, byte_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { + store_with_pointer_offset(frag, 0); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIteratorResidualLast for interleaved data. +/// It is mapped to the congruous layout. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// + +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int AccessSize, + int InterleavedK> +class PredicatedTileIteratorResidualLast< + Shape_, + Element_, + layout::ColumnMajorInterleaved, + AdvanceRank, + ThreadMap_, + AccessSize, + false> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + static int const kInterleavedK = InterleavedK; + using Layout = layout::ColumnMajorInterleaved; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingIterator = PredicatedTileIteratorResidualLast< + layout::PitchLinearShape< + Shape::kRow * kInterleavedK, + Shape::kColumn / kInterleavedK>, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 0 : 1), + ThreadMap, + AccessSize>; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = cutlass::Array< + Element, + ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) + : params_(layout::PitchLinear(layout.stride(0))) {} + + CUTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const& base) + : params_(base) {} + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const& threadblock_offset, + int const* indices = + nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : iterator_( + params.params_, + pointer, + layout::PitchLinearCoord( + extent.row() * kInterleavedK, + extent.column() / kInterleavedK), + thread_id, + layout::PitchLinearCoord( + threadblock_offset.row() * kInterleavedK, + threadblock_offset.column() / kInterleavedK)) {} + + /// Construct a PredicatedTileIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIteratorResidualLast( + params, + pointer, + extent, + thread_id, + make_Coord(0, 0)) {} + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast& operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast operator++(int) { + PredicatedTileIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + iterator_.clear_mask(enable); + } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { + iterator_.set_residual_tile(enable); + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { + iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { + iterator_.set_mask(mask); + } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { + iterator_.get_mask(mask); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment& frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) { + load_with_pointer_offset(frag, 0); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { + store_with_pointer_offset(frag, 0); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIteratorResidualLast for interleaved-32 +/// data. It is mapped to the congruous layout. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int AccessSize, + int InterleavedK> +class PredicatedTileIteratorResidualLast< + Shape_, + Element_, + layout::RowMajorInterleaved, + AdvanceRank, + ThreadMap_, + AccessSize, + false> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + static int const kInterleavedK = InterleavedK; + using Layout = layout::RowMajorInterleaved; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingIterator = PredicatedTileIteratorResidualLast< + layout::PitchLinearShape< + Shape::kColumn * kInterleavedK, + Shape::kRow / kInterleavedK>, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap, + AccessSize>; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = cutlass::Array< + Element, + ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) + : params_(layout::PitchLinear(layout.stride(0))) {} + + CUTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const& base) + : params_(base) {} + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const& threadblock_offset, + int const* indices = + nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : iterator_( + params.params_, + pointer, + layout::PitchLinearCoord( + extent.column() * kInterleavedK, + extent.row() / kInterleavedK), + thread_id, + layout::PitchLinearCoord( + threadblock_offset.column() * kInterleavedK, + threadblock_offset.row() / kInterleavedK)) {} + + /// Construct a PredicatedTileIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIteratorResidualLast( + params, + pointer, + extent, + thread_id, + make_Coord(0, 0)) {} + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast& operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast operator++(int) { + PredicatedTileIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + iterator_.clear_mask(enable); + } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { + iterator_.set_residual_tile(enable); + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { + iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { + iterator_.set_mask(mask); + } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { + iterator_.get_mask(mask); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment& frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) { + load_with_pointer_offset(frag, 0); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { + store_with_pointer_offset(frag, 0); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fmha/iterators/transpose_warp_iterator.h b/natten/csrc/include/natten/cuda/fmha/iterators/transpose_warp_iterator.h new file mode 100644 index 00000000..4e9faa96 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha/iterators/transpose_warp_iterator.h @@ -0,0 +1,60 @@ +/* + * Copied from xFormers (https://github.com/facebookresearch/xformers/) + * + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * BSD 3-Clause License + * + * 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 names of Facebook, Deepmind Technologies, NYU, NEC + * Laboratories America and IDIAP Research Institute 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 OWNER 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 + +template +struct TransposeWarpIterator { + using Iterator = char; + static bool constexpr kSupportsTranspose = false; +}; + +template < + /// Operand identity + cutlass::gemm::Operand Operand, + /// Data type of A elements + typename Element, + typename InstructionShape, + bool kTranspose> +struct TransposeWarpIterator< + cutlass::gemm::warp:: + WarpIteratorFromSmem> { + using Iterator = cutlass::gemm::warp:: + WarpIteratorFromSmem; + static bool constexpr kSupportsTranspose = true; +}; diff --git a/natten/csrc/include/natten/cuda/fmha/iterators/warp_iterator_from_smem.h b/natten/csrc/include/natten/cuda/fmha/iterators/warp_iterator_from_smem.h new file mode 100644 index 00000000..989148d0 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha/iterators/warp_iterator_from_smem.h @@ -0,0 +1,290 @@ +/* + * Copied from xFormers (https://github.com/facebookresearch/xformers/). + * + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + */ +/*************************************************************************************************** + * Copyright (c) 2017 - 2026 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 Inspired from + "cutlass/gemm/warp/mma_tensor_op_tile_access_iterator.h" Loads tiles of GEMM + operands from a RowMajor shared-memory layout into registers to use by A100 + TensorCores. + + The difference with "mma_tensor_op_tile_access_iterator.h" is that: + (1) We use "ldmatrix" to load tiles, rather than manual loads (slightly + faster) (2) We support to transpose the operand (eg read `A.transpose()` when + the shared memory holds `A`) + + This is only implemented for the specific shapes. +*/ +#pragma once + +#include + +//////////////////////////////////////////////////////////////////////////////// +namespace cutlass { +namespace gemm { +namespace warp { + +template < + /// Operand identity + Operand Operand_, + /// Data type of A elements + typename Element_, + typename InstructionShape_, + bool kTranspose = false> +class WarpIteratorFromSmem { + public: + /// Shape of tile to load (concept: MatrixShape) + using Shape = cutlass::MatrixShape<32, 32>; + + /// Operand tag + static Operand const kOperand = Operand_; + static_assert( + kOperand == Operand::kA, + "No support for OperandB at the moment"); + + /// Basic check + static_assert( + kOperand == Operand::kA || kOperand == Operand::kB, + "WarpIteratorFromSmem may only be instantiated for A or B operands to warp-level Mma."); + + /// Element type + using Element = Element_; + static_assert(sizeof_bits::value == 16, "Only supported for half"); + + /// Layout of source tile + using Layout = cutlass::layout::RowMajor; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + static_assert(InstructionShape::kRow == 16, "Only supports 16x8x8 / 16x8x16"); + static_assert( + InstructionShape::kColumn == 8 || InstructionShape::kColumn == 16, + "Only supports 16x8x8 / 16x8x16"); + + /// Delta between *MMA operations (in units of *MMA operations, concept: + /// MatrixShape) + static int const kOpDelta = 1; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Number of elements accessed per Shared Memory load + static int const kElementsPerAccess = + (sizeof_bits::value >= 32 ? 1 + : 32 / sizeof_bits::value); + + using InstructionCount = MatrixShape< + Shape::kRow / InstructionShape::kRow, + Shape::kColumn / InstructionShape::kColumn>; + + static int const kIterations = (kOperand == Operand::kA) + ? InstructionCount::kColumn + : InstructionCount::kRow; + + public: + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = Array< + Element, + (kOperand == Operand::kA) + ? (Shape::kRow* InstructionShape::kColumn / kThreads) + : (Shape::kColumn* InstructionShape::kRow / kThreads)>; + + /// Memory access type + // using AccessType = AlignedArray; + using AccessType = Array; + + static int constexpr kWarpShapeDivisibleInner = + (kOperand == Operand::kA ? InstructionShape::kColumn + : InstructionShape::kRow); + static int constexpr kAccessesInner = + (kWarpShapeDivisibleInner / kElementsPerAccess) / 4; + // Number of 32bits tiles to load per `ldmatrix` + static int const kTilesPerInstruction = InstructionShape::kRow / 8; + static_assert(kTilesPerInstruction == 2, "Only supports 16x8x16 and 16x8x8"); + + private: + /// Underlying tensor reference + TensorRef ref_; + + /// Origin + MatrixCoord origin_; + + /// Iterations in a tile + int iterations_; + + public: + /// Constructor from TensorRef + CUTLASS_HOST_DEVICE + WarpIteratorFromSmem(TensorRef const& ref, int lane_id) + : WarpIteratorFromSmem(ref, {Shape::kRow, Shape::kColumn}, lane_id) {} + CUTLASS_HOST_DEVICE + WarpIteratorFromSmem(TensorRef const& ref, TensorCoord extent, int lane_id) + : ref_(ref), iterations_(0) { + // See also: + // https://docs.nvidia.com/cuda/archive/11.7.1/parallel-thread-execution/index.html#warp-level-matrix-fragment-mma-1688 + // 16x8x8: kAccessesInner = 1 (1 ldmatrix.x4) + // 16x8x16: kAccessesInner = 2 (2 ldmatrix.x4) + int ldsm_vec_num = (lane_id >> 3); + if (kOperand == Operand::kA) { + origin_ = MatrixCoord(lane_id % 8, 0); + static_assert( + InstructionCount::kRow * kTilesPerInstruction == 4, + "can't use ldmatrix.x4"); + int access_m_idx = ldsm_vec_num % kTilesPerInstruction; + int inner_idx = (ldsm_vec_num / kTilesPerInstruction) % kAccessesInner; + int inst_m_idx = ldsm_vec_num / (kTilesPerInstruction * kAccessesInner); + MatrixCoord offset( + access_m_idx * 8 + inst_m_idx * InstructionShape::kRow, + inner_idx * 4 * kElementsPerAccess); + if (kTranspose) { + offset = MatrixCoord(offset.column(), offset.row()); + } + origin_ += offset; + } else { + // XXX: This is not tested or used + origin_ = MatrixCoord(0, lane_id % 8); + static_assert(InstructionCount::kColumn * kAccessesInner == 4, ""); + CUTLASS_PRAGMA_UNROLL + for (int inst_n_idx = 0; inst_n_idx < InstructionCount::kColumn; + ++inst_n_idx) { + CUTLASS_PRAGMA_UNROLL + for (int inner_idx = 0; inner_idx < kAccessesInner; ++inner_idx) { + int access_idx = inner_idx + kAccessesInner * inst_n_idx; + + MatrixCoord offset( + inner_idx * 4 * kElementsPerAccess, inst_n_idx * 8); + + if (access_idx == ldsm_vec_num) { + if (kTranspose) { + offset = MatrixCoord(offset.column(), offset.row()); + } + origin_ += offset; + } + } + } + } + + ref_.add_coord_offset(origin_); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_HOST_DEVICE + WarpIteratorFromSmem& add_tile_offset(TensorCoord const& tile_offset) { + TensorCoord coord_offset( + tile_offset.row() * Shape::kRow, tile_offset.column() * Shape::kColumn); + if (kTranspose) { + coord_offset = TensorCoord{coord_offset.column(), coord_offset.row()}; + } + origin_ += coord_offset; + + ref_.add_coord_offset(coord_offset); + + return *this; + } + + /// Advances the iterator along the advance dimension + CUTLASS_DEVICE + void advance() { + if (kOperand == Operand::kA) { + add_tile_offset({0, 1}); + } else { + add_tile_offset({1, 0}); + } + + iterations_ = 0; + } + + /// increase iterations in a tile + CUTLASS_HOST_DEVICE + WarpIteratorFromSmem& operator++() { + iterations_++; + + if (iterations_ >= kIterations) + advance(); + + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + CUTLASS_DEVICE + void load(Fragment& frag) const { + AccessType* access_ptr = reinterpret_cast(&frag); + using LoadLayout = typename platform:: + conditional::type; + + CUTLASS_PRAGMA_UNROLL + for (int access_m_idx = 0; access_m_idx < + (InstructionCount::kRow * kTilesPerInstruction * kAccessesInner) / 4; + ++access_m_idx) { + MatrixCoord offset; + if (kOperand == Operand::kA) { + offset = MatrixCoord( + access_m_idx * 16, iterations_ * InstructionShape::kColumn); + } else { + offset = MatrixCoord(iterations_ * InstructionShape::kRow, 0); + } + if (kTranspose) { + offset = MatrixCoord(offset.column(), offset.row()); + } + cutlass::arch::ldsm( + access_ptr[access_m_idx], ref_.data() + ref_.offset(offset)); + } + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace cutlass +//////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fmha/kernel_backward.h b/natten/csrc/include/natten/cuda/fmha/kernel_backward.h new file mode 100644 index 00000000..e548d1d5 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha/kernel_backward.h @@ -0,0 +1,2282 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ +#pragma once + +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include +#include + +#include + +namespace natten { +namespace cuda { +namespace fmha { + +using namespace gemm_kernel_utils; + +namespace { + +template +struct GmemTile { + /* + Helper functions to efficient store/load RF to gmem + + GEMM accumulators have a particular format on A100, and + it takes some compute/shared-memory to rearrange them to + a RowMajor or ColumnMajor format in global memory through + an Epilogue. The same complexity goes for loading into RF. + + This class loads/stores RF as they are, and can be used for + efficient accumulation across gemms for instance: + + ``` + GmemTile tile; + for (int i = 0; i < N; ++i) { + // ... + + Fragment accum; + if (i == 0) { + accum.clear(); + } else { + tile.load(accum); + } + mma(accum, ...); + if (i < N-1) { + // Store for next GEMM + tile.store(accum); + } else { + // Store in tensor (eg RowMajor) + epilogue(accum); + } + + // ... + } + ``` + */ + + // 128bits per thread + using AccessType = cutlass::Array; + static constexpr int32_t kBytes = sizeof(AccessType); + static constexpr int32_t kStride = kNumThreads * AccessType::kElements; + static constexpr int32_t kNumIters = + FragmentType::kElements / AccessType::kElements; + static constexpr int32_t kElementsStored = + kNumThreads * FragmentType::kElements; + static_assert( + FragmentType::kElements % AccessType::kElements == 0, + "fragment not aligned on 128 bits"); + + float* ptr; + + CUTLASS_DEVICE void load(FragmentType& fragment, int thread_id) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kNumIters; ++i) { + AccessType* __restrict__ gmem_ptr = reinterpret_cast( + ptr + thread_id * AccessType::kElements + i * kStride); + AccessType sub_fragment; + cutlass::arch::global_load( + sub_fragment, gmem_ptr, true); + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < AccessType::kElements; ++j) { + fragment[i * AccessType::kElements + j] = sub_fragment[j]; + } + } + } + + CUTLASS_DEVICE void store(FragmentType const& fragment, int thread_id) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kNumIters; ++i) { + AccessType* __restrict__ gmem_ptr = reinterpret_cast( + ptr + thread_id * AccessType::kElements + i * kStride); + AccessType sub_fragment; + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < AccessType::kElements; ++j) { + sub_fragment[j] = fragment[i * AccessType::kElements + j]; + } + cutlass::arch::global_store( + sub_fragment, gmem_ptr, true); + } + } +}; + +struct AtomicLock { + CUTLASS_DEVICE static void acquire( + int32_t* lock, + int set_val, + int thread_id) { + if (thread_id == 0) { + while (atomicCAS(lock, 0 /*cmp*/, set_val /*setval*/) != set_val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 + __nanosleep(40); +#endif + } + } + __syncthreads(); + } + CUTLASS_DEVICE static void release(int32_t* lock, int thread_id) { + if (thread_id == 0) { + int status = 0; +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 + asm volatile("st.global.release.gpu.b32 [%0], %1;\n" + : + : "l"(lock), "r"(status)); +#else + asm volatile("st.global.cg.b32 [%0], %1;\n" : : "l"(lock), "r"(status)); +#endif + } + } +}; + +template +constexpr int getWarpsPerSmBw() { + bool is_half = !cutlass::platform::is_same::value; + if (Arch::kMinComputeCapability >= 80) { + return is_half ? 12 : 8; + } + return 8; +} +} // namespace + +template < + // input/output type + typename scalar_t_, + // which arch we target (eg `cutlass::arch::Sm80`) + typename ArchTag_, + // run optimized kernel because memory accesses will be aligned + bool kIsAligned_, + // NOTE(alih): we force alignment on dim in NATTEN, so we + // only need kPreload_=true instantiations. + //// when doing a GEMM, preload the next one (uses more shmem) + //// bool kPreload_ = true, + // block dimensions + int kBlockSizeI_, + int kBlockSizeJ_, + // upperbound on `max(value.shape[-1], query.shape[-1])` + int kMaxK_ = (int)cutlass::platform::numeric_limits::max(), + // assumes that `cu_seqlen` is None, and + // (1) `num_queries % kBlockSizeI == 0` + // (2) `num_keys % kBlockSizeJ == 0` + bool kKeysQueriesAlignedToBlockSize_ = false> +struct AttentionBackwardKernel { + enum CustomMaskType { + NoCustomMask = 0, + CausalFromTopLeft = 1, + CausalFromBottomRight = 2, + NumCustomMaskTypes, + }; + using scalar_t = scalar_t_; + using output_t = scalar_t; + using output_accum_t = float; + using lse_scalar_t = float; + using accum_t = float; + using ArchTag = ArchTag_; + static constexpr bool kIsAligned = kIsAligned_; + // static constexpr bool kPreload = kPreload_; + static constexpr int kBlockSizeI = kBlockSizeI_; + static constexpr int kBlockSizeJ = kBlockSizeJ_; + static constexpr int kMaxK = kMaxK_; + static constexpr bool kKeysQueriesAlignedToBlockSize = + kKeysQueriesAlignedToBlockSize_; + + static constexpr int64_t kWarpSize = 32; + + // If this is true, we store and accumulate dK/dV in RF + // rather than going back to gmem everytime + static constexpr bool kIsHalf = cutlass::sizeof_bits::value <= 16; + static constexpr bool kOutputInRF = kIsHalf && kMaxK <= kBlockSizeI; + + static constexpr bool kPreload = + kIsHalf && ArchTag::kMinComputeCapability >= 80 && kOutputInRF; + static_assert( + !kPreload || + (kIsHalf && ArchTag::kMinComputeCapability >= 80 && kOutputInRF), + "preload MMA not supported"); + static constexpr bool kPrologueQK = kPreload; + static constexpr bool kPrologueGV = kPreload; + static constexpr bool kPrologueDOV = kPreload; + static constexpr bool kPrologueGQ = kPreload; + static constexpr bool kPrologueGK = kPreload; + + static constexpr int64_t kNumWarpsPerBlock = + (kBlockSizeI * kBlockSizeJ) / (32 * 32); + + // Compute delta for the f16 kernels + // TODO: Figure out why it's slower on the f32 kernels + // (something due to RF pressure?) + // TODO: Remove condition on `kOutputInRF` - this is needed to work + // around a compiler bug on V100, not exactly sure why but I spent + // too much time on this already. Reproducible with + // (B, Mq, Mkv, K) = (1, 1, 1, 136) for instance + // TODO (alih): this was buggy when I was doing FNA in 2024, so disabling for + // safety + // static constexpr bool kKernelComputesDelta = + // kIsHalf && (kOutputInRF || ArchTag::kMinComputeCapability != 70); + static constexpr bool kKernelComputesDelta = false; + + // Launch bounds + static constexpr int64_t kNumThreads = kWarpSize * kNumWarpsPerBlock; + static constexpr int64_t kMinBlocksPerSm = + getWarpsPerSmBw() / kNumWarpsPerBlock; + + using GemmType = DefaultGemmType; + using DefaultConfig = + typename cutlass::gemm::device::DefaultGemmConfiguration< + typename GemmType::OpClass, + ArchTag, + scalar_t, + scalar_t, + scalar_t, // ElementC + accum_t // ElementAccumulator + >; + static constexpr auto kOptimalAlignement = cutlass::platform::max( + DefaultConfig::kAlignmentA, + DefaultConfig::kAlignmentB); + static constexpr auto kMinimumAlignment = GemmType::kMinimumAlignment; + + struct MatmulQK { + /* + attn_T = k_j @ q_i.transpose(-2, -1) # matmul + attn_T = (attn_T - logsumexp[i_start:i_end].unsqueeze(1).transpose(-2, + -1)).exp() # epilogue + + with attn_T.shape = (kBlockSizeJ, kBlockSizeI) + */ + using ThreadblockShape = + cutlass::gemm::GemmShape; + using WarpShape = cutlass::gemm::GemmShape<32, 32, GemmType::WarpK>; + using DefaultMma = typename cutlass::gemm::threadblock::DefaultMma< + scalar_t, // ElementA + cutlass::layout::RowMajor, // LayoutA + kIsAligned ? DefaultConfig::kAlignmentA : GemmType::kMinimumAlignment, + scalar_t, // ElementB + cutlass::layout::ColumnMajor, // LayoutB + kIsAligned ? DefaultConfig::kAlignmentB : GemmType::kMinimumAlignment, + accum_t, // ElementC + cutlass::layout::RowMajor, // LayoutC + typename GemmType::OpClass, + ArchTag, + ThreadblockShape, + WarpShape, + typename GemmType::InstructionShape, + DefaultConfig::kStages, + typename GemmType::Operator, + false, // AccumulatorsInRowMajor = false, + cutlass::gemm::SharedMemoryClearOption::kNone>; + using MmaCore = typename DefaultMma::MmaCore; + using Mma = + typename MakeCustomMma::Mma; + + // Epilogue to store to shared-memory in a format that we can use later for + // the second matmul + using B2bGemm = typename cutlass::gemm::threadblock::B2bGemm< + typename Mma::Operator::IteratorC, + typename Mma::Operator, + scalar_t, + WarpShape, + ThreadblockShape>; + using AccumLambdaIterator = typename DefaultMmaAccumLambdaIterator< + typename Mma::Operator::IteratorC, + accum_t, + kWarpSize>::Iterator; + using AccumulatorSharedStorage = typename B2bGemm::AccumulatorSharedStorage; + }; + + struct MatmulGradV { + /* + grad_v[j_start:j_end] += attn_T @ do_i # matmul + + Dimensions: (kBlockSizeJ * kNumWarpsPerBlock, kBlockSizeI, K) + (we might need to iterate multiple times on K) + */ + using ThreadblockShape = + cutlass::gemm::GemmShape; + using WarpShape = cutlass::gemm::GemmShape<32, 32, GemmType::WarpK>; + using InstructionShape = typename GemmType::InstructionShape; + + using DefaultGemm = cutlass::gemm::kernel::DefaultGemm< + scalar_t, // ElementA, + cutlass::layout::RowMajor, // LayoutA, + DefaultConfig::kAlignmentA, + scalar_t, // ElementB, + cutlass::layout::RowMajor, // LayoutB, + kIsAligned ? DefaultConfig::kAlignmentB : GemmType::kMinimumAlignment, + output_t, + cutlass::layout::RowMajor, // LayoutC, + accum_t, + typename GemmType::OpClass, + ArchTag, + ThreadblockShape, + WarpShape, + typename GemmType::InstructionShape, + typename DefaultConfig::EpilogueOutputOp, + void, // ThreadblockSwizzle - not used + DefaultConfig::kStages, + false, // SplitKSerial + typename GemmType::Operator>; + + using WarpIteratorA = typename cutlass::gemm::threadblock:: + DefaultWarpIteratorAFromSharedMemory< + typename DefaultGemm::Mma::Operator::Shape, // WarpShape + typename DefaultGemm::Mma::Operator:: + InstructionShape, // InstructionShape + typename DefaultGemm::Mma::Operator:: + IteratorA, // RegularWarpIterator + typename DefaultGemm::Mma::Policy // Policy + >::WarpIterator; + using DefaultMmaFromSmem = + typename cutlass::gemm::threadblock::DefaultMmaFromSharedMemory< + typename DefaultGemm::Mma, + MatmulQK::AccumulatorSharedStorage::Shape::kN, + WarpIteratorA, + /*kApplyDropout*/ false>; // kScaleOperandA + + using Mma = typename DefaultMmaFromSmem::Mma; + using IteratorB = typename Mma::IteratorB; + using WarpCount = typename Mma::WarpCount; + + // Epilogue + using DefaultOutputOp = typename DefaultConfig::EpilogueOutputOp; + using DefaultEpilogue = typename DefaultGemm::Epilogue; + using OutputTileIterator = + typename cutlass::epilogue::threadblock::MakePrefetchableIterator< + typename DefaultEpilogue::OutputTileIterator>::Iterator; + using AccumTileGmem = GmemTile; + }; + + struct MatmulDOIVJ { + /* + doi_t_vj = do_i @ v_j.transpose(-2, -1) # matmul + tmp = (doi_t_vj - Di.unsqueeze(1)) * attn # inplace / epilogue? + */ + using ThreadblockShape = + cutlass::gemm::GemmShape; + using WarpShape = cutlass::gemm::GemmShape<32, 32, GemmType::WarpK>; + + using ElementC = output_t; + using ElementAccum = accum_t; + + // no-op output op - epilogue just stores result to global memory + using BiasGradEpilogueOutputOp = + typename cutlass::epilogue::thread::LinearCombination< + ElementC, + DefaultConfig::EpilogueOutputOp::kCount, + typename DefaultConfig::EpilogueOutputOp::ElementAccumulator, + typename DefaultConfig::EpilogueOutputOp::ElementCompute, + cutlass::epilogue::thread::ScaleType::Nothing>; + + using DefaultGemm = typename cutlass::gemm::kernel::DefaultGemm< + scalar_t, // ElementA + cutlass::layout::RowMajor, // LayoutA + kIsAligned ? DefaultConfig::kAlignmentA : GemmType::kMinimumAlignment, + scalar_t, // ElementB + cutlass::layout::ColumnMajor, // LayoutB + kIsAligned ? DefaultConfig::kAlignmentB : GemmType::kMinimumAlignment, + ElementC, // ElementC + cutlass::layout::RowMajor, // LayoutC + ElementAccum, // ElementAccumulator + typename GemmType::OpClass, + ArchTag, + ThreadblockShape, + WarpShape, + typename GemmType::InstructionShape, + BiasGradEpilogueOutputOp, // EpilogueOutputOp + void, // ThreadblockSwizzle (not used) + /*// multiple preloads, and 3 stages push us over shared + // memory capacity on A100. + kPreload && (kBlockSizeI * kBlockSizeJ > 64 * 64) + ? cutlass::const_min(2, DefaultConfig::kStages) + : */ + DefaultConfig::kStages, // Stages + false, // SplitKSerial + typename GemmType::Operator, + cutlass::gemm::SharedMemoryClearOption::kNone>; + using Mma = typename MakeCustomMma::Mma; + using AccumLambdaIterator = typename DefaultMmaAccumLambdaIterator< + typename Mma::Operator::IteratorC, + ElementAccum, + kWarpSize>::Iterator; + + // epilogue used to write bias gradient, which is just the output of this + // matmul with some operations applied to the fragment + using BiasGradEpilogue = typename DefaultGemm::Epilogue; + + // Epilogue to store to shared-memory in a format that we can use later for + // the second matmul + using B2bGemm = typename cutlass::gemm::threadblock::B2bGemm< + typename DefaultGemm::Mma::Operator::IteratorC, + typename DefaultGemm::Mma::Operator, + scalar_t, + WarpShape, + ThreadblockShape>; + using AccumulatorSharedStorage = typename B2bGemm::AccumulatorSharedStorage; + }; + + struct MatmulGradQ { + // grad_q <- tmp @ k_j + using ThreadblockShape = + cutlass::gemm::GemmShape; + using WarpShape = cutlass::gemm::GemmShape<32, 32, GemmType::WarpK>; + using InstructionShape = typename GemmType::InstructionShape; + + using DefaultGemm = cutlass::gemm::kernel::DefaultGemm< + scalar_t, // ElementA, + cutlass::layout::RowMajor, // LayoutA, + DefaultConfig::kAlignmentA, + scalar_t, // ElementB, + cutlass::layout::RowMajor, // LayoutB, + kIsAligned ? DefaultConfig::kAlignmentB : GemmType::kMinimumAlignment, + output_t, + cutlass::layout::RowMajor, // LayoutC, + accum_t, + typename GemmType::OpClass, + ArchTag, + ThreadblockShape, + WarpShape, + typename GemmType::InstructionShape, + typename DefaultConfig::EpilogueOutputOp, + void, // ThreadblockSwizzle - not used + DefaultConfig::kStages, + false, // SplitKSerial + typename GemmType::Operator>; + + using WarpIteratorA = typename cutlass::gemm::threadblock:: + DefaultWarpIteratorAFromSharedMemory< + typename DefaultGemm::Mma::Operator::Shape, + typename DefaultGemm::Mma::Operator::InstructionShape, + typename DefaultGemm::Mma::Operator::IteratorA, + typename DefaultGemm::Mma::Policy>::WarpIterator; + using DefaultMmaFromSmem = + typename cutlass::gemm::threadblock::DefaultMmaFromSharedMemory< + typename DefaultGemm::Mma, + MatmulDOIVJ::AccumulatorSharedStorage::Shape::kN, + WarpIteratorA, + false>; // kScaleOperandA + using Mma = typename DefaultMmaFromSmem::Mma; + using IteratorB = typename Mma::IteratorB; + using WarpCount = typename Mma::WarpCount; + + // Epilogue + using DefaultOutputOp = typename DefaultConfig::EpilogueOutputOp; + using DefaultEpilogue = typename DefaultGemm::Epilogue; + using OutputTileIterator = + typename cutlass::epilogue::threadblock::MakePrefetchableIterator< + typename DefaultEpilogue::OutputTileIterator>::Iterator; + using AccumTileGmem = GmemTile; + }; + + struct MatmulGradK { + // grad_k <- tmp.transpose(-2, -1) @ q_i + using ThreadblockShape = + cutlass::gemm::GemmShape; + using WarpShape = cutlass::gemm::GemmShape<32, 32, GemmType::WarpK>; + using InstructionShape = typename GemmType::InstructionShape; + + using DefaultGemm = cutlass::gemm::kernel::DefaultGemm< + scalar_t, // ElementA, + cutlass::layout::RowMajor, // LayoutA, + DefaultConfig::kAlignmentA, + scalar_t, // ElementB, + cutlass::layout::RowMajor, // LayoutB, + kIsAligned ? DefaultConfig::kAlignmentB : GemmType::kMinimumAlignment, + output_t, + cutlass::layout::RowMajor, // LayoutC, + accum_t, + typename GemmType::OpClass, + ArchTag, + ThreadblockShape, + WarpShape, + typename GemmType::InstructionShape, + typename DefaultConfig::EpilogueOutputOp, + void, // ThreadblockSwizzle - not used + DefaultConfig::kStages, + false, // SplitKSerial + typename GemmType::Operator>; + + using WarpIteratorA = typename cutlass::gemm::threadblock:: + DefaultWarpIteratorAFromSharedMemory< + typename DefaultGemm::Mma::Operator::Shape, + typename DefaultGemm::Mma::Operator::InstructionShape, + typename DefaultGemm::Mma::Operator::IteratorA, + typename DefaultGemm::Mma::Policy>::WarpIterator; + using DefaultMmaFromSmemN = + typename cutlass::gemm::threadblock::DefaultMmaFromSharedMemory< + typename DefaultGemm::Mma, + MatmulQK::AccumulatorSharedStorage::Shape::kN, // kMaxK + WarpIteratorA, + false>; // kScaleOperandA + using DefaultMmaFromSmemT = + typename cutlass::gemm::threadblock::DefaultMmaFromSharedMemory< + typename DefaultGemm::Mma, + MatmulDOIVJ::AccumulatorSharedStorage::Shape::kM, // kMaxK + WarpIteratorA, + false, // kScaleOperandA + kPreload>; // kTransposeA + using DefaultMmaFromSmem = typename cutlass::platform::conditional< + DefaultMmaFromSmemT::kIsTransposedA, + DefaultMmaFromSmemT, + DefaultMmaFromSmemN>::type; + using Mma = typename DefaultMmaFromSmem::Mma; + using IteratorB = typename Mma::IteratorB; + using WarpCount = typename Mma::WarpCount; + + // Epilogue + using DefaultOutputOp = typename DefaultConfig::EpilogueOutputOp; + using DefaultEpilogue = typename DefaultGemm::Epilogue; + using OutputTileIterator = + typename cutlass::epilogue::threadblock::MakePrefetchableIterator< + typename DefaultEpilogue::OutputTileIterator>::Iterator; + using AccumTileGmem = GmemTile; + }; + + static constexpr bool kEnableSplitKeys = true; + + static constexpr bool kNeedsAccumGradQ = kEnableSplitKeys || + !cutlass::platform::is_same::value; + static constexpr bool kNeedsAccumGradK = !kOutputInRF && + !cutlass::platform::is_same::value; + static constexpr bool kNeedsAccumGradV = !kOutputInRF && + !cutlass::platform::is_same::value; + + struct GradQTempStorage { + int32_t lock; + int32_t counter; + int32_t pad[2]; // pad to 128bits + output_accum_t buffer[MatmulGradQ::AccumTileGmem::kElementsStored]; + }; + + struct Params { + // Input tensors + scalar_t* query_ptr = nullptr; // [Mq, nH, K] + scalar_t* key_ptr = nullptr; // [Mk, nH, K] + scalar_t* value_ptr = nullptr; // [Mk, nH, Kv] + lse_scalar_t* logsumexp_ptr = nullptr; // [Mq, nH] + scalar_t* output_ptr = nullptr; // [Mq, nH, Kv] + scalar_t* grad_output_ptr = nullptr; // [Mq, nH, Kv] + accum_t* delta_ptr = nullptr; // [nH, Mq] + + const int32_t* ptr_cumulative_seqlen_Q = nullptr; + const int32_t* ptr_cumulative_seqlen_KV = nullptr; + + // Output tensors + output_t* grad_query_ptr = nullptr; // [Mq, nH, K] + output_t* grad_key_ptr = nullptr; // [Mk, nH, K] + output_t* grad_value_ptr = nullptr; // [Mk, nH, Kv] + + // Accumulators + output_accum_t* workspace = nullptr; // [Mq, Kq] + [Mkv, Kq] + [Mkv, Kv] + output_accum_t* workspace_gv = + nullptr; // (will be calculated by the kernel) + GradQTempStorage* workspace_gq = + nullptr; // (will be calculated by the kernel) + + // Scale + accum_t scale = 1.0f; + + // Dimensions/strides + int32_t head_dim = -1; + int32_t head_dim_value = -1; + int32_t num_queries = -1; + int32_t num_keys = -1; + int32_t num_heads = -1; + uint8_t custom_mask_type = NoCustomMask; + + int32_t q_strideM = -1; + int32_t k_strideM = -1; + int32_t v_strideM = -1; + int32_t gO_strideM = -1; + + CUTLASS_HOST_DEVICE int32_t o_strideM() const { + return head_dim_value * num_heads; + } + CUTLASS_HOST_DEVICE int32_t gQ_strideM() const { + return num_heads * head_dim; + } + CUTLASS_HOST_DEVICE int32_t gK_strideM() const { + return num_heads * head_dim; + } + CUTLASS_HOST_DEVICE int32_t gV_strideM() const { + return num_heads * head_dim_value; + } + + // Everything below is only used in `advance_to_block` + // and shouldn't use registers + int64_t o_strideH = -1; + int32_t q_strideH = -1; + int32_t k_strideH = -1; + int32_t v_strideH = -1; + int64_t o_strideB = -1; + int64_t q_strideB = -1; + int64_t k_strideB = -1; + int64_t v_strideB = -1; + int64_t lse_strideB = -1; + int64_t delta_strideB = -1; + int32_t num_batches = -1; + int16_t num_splits_key = 1; // We use `gridDim.x` inside kernel + + int64_t gO_strideB = 0; + int64_t gQ_strideB = 0; + int64_t gK_strideB = 0; + int64_t gV_strideB = 0; + int64_t gO_strideH = 0; + int64_t gQ_strideH = 0; + int64_t gK_strideH = 0; + int64_t gV_strideH = 0; + + CUTLASS_HOST_DEVICE int16_t num_splits_key_device() const { +#ifdef __CUDA_ARCH__ + return kEnableSplitKeys ? gridDim.x : 1; +#else + return num_splits_key; // for host-side tests +#endif + } + + CUTLASS_HOST_DEVICE int16_t split_key_device() const { +#ifdef __CUDA_ARCH__ + return kEnableSplitKeys ? blockIdx.x : 0; +#else + return 0; // for host-side tests +#endif + } + + CUTLASS_DEVICE bool advance_to_block() { + int64_t batch_id = blockIdx.z; + int32_t head_id = blockIdx.y; + + if (kNeedsAccumGradQ || kNeedsAccumGradK || kNeedsAccumGradV) { + assert(workspace_size() == 0 || workspace != nullptr); + + workspace += (batch_id * num_heads + head_id) * workspace_strideBH(); + workspace = warp_uniform(workspace); + workspace_gv = workspace + workspace_elements_gk(); + workspace_gq = + (GradQTempStorage*)(workspace_gv + workspace_elements_gv()); + if (kEnableSplitKeys) { + workspace_gv += workspace_elements_gv() * split_key_device() / + num_splits_key_device(); + workspace += workspace_elements_gk() * split_key_device() / + num_splits_key_device(); + } + } else { + workspace = nullptr; + } + + // Advance pointers that depend on the total concatenated + // number of queries, as `num_queries` is modified in the block + // below + + if (ptr_cumulative_seqlen_Q != nullptr && + ptr_cumulative_seqlen_KV != nullptr) { + ptr_cumulative_seqlen_Q += batch_id; + ptr_cumulative_seqlen_KV += batch_id; + int32_t q_start = ptr_cumulative_seqlen_Q[0]; + int32_t k_start = ptr_cumulative_seqlen_KV[0]; + int32_t q_next_start = ptr_cumulative_seqlen_Q[1]; + int32_t k_next_start = ptr_cumulative_seqlen_KV[1]; + assert(q_next_start - q_start <= num_queries); + assert(k_next_start - k_start <= num_keys); + num_queries = q_next_start - q_start; + num_keys = k_next_start - k_start; + + // Jump manually + batch_id = 0; + + query_ptr += q_start * q_strideM; + key_ptr += k_start * k_strideM; + value_ptr += k_start * v_strideM; + output_ptr += q_start * o_strideM(); + grad_output_ptr += q_start * gO_strideM; + delta_ptr += q_start * num_heads; + logsumexp_ptr += q_start * num_heads; + + grad_query_ptr += q_start * gQ_strideM(); + grad_key_ptr += k_start * gK_strideM(); + grad_value_ptr += k_start * gV_strideM(); + } + + logsumexp_ptr += batch_id * lse_strideB + head_id; + delta_ptr += batch_id * delta_strideB + head_id; + + query_ptr += batch_id * q_strideB + head_id * q_strideH; + key_ptr += batch_id * k_strideB + head_id * k_strideH; + value_ptr += batch_id * v_strideB + head_id * v_strideH; + output_ptr += batch_id * o_strideB + head_id * o_strideH; + grad_output_ptr += batch_id * gO_strideB + head_id * gO_strideH; + + grad_query_ptr += batch_id * gQ_strideB + head_id * gQ_strideH; + grad_key_ptr += batch_id * gK_strideB + head_id * gK_strideH; + grad_value_ptr += batch_id * gV_strideB + head_id * gV_strideH; + + // Some values are modified above + // Signal to the compiler that they are the same in all threads + // and can be stored in warp-uniform registers (Sm75+) + num_queries = warp_uniform(num_queries); + num_keys = warp_uniform(num_keys); + custom_mask_type = warp_uniform(custom_mask_type); + + query_ptr = warp_uniform(query_ptr); + key_ptr = warp_uniform(key_ptr); + value_ptr = warp_uniform(value_ptr); + logsumexp_ptr = warp_uniform(logsumexp_ptr); + output_ptr = warp_uniform(output_ptr); + grad_output_ptr = warp_uniform(grad_output_ptr); + delta_ptr = warp_uniform(delta_ptr); + + grad_query_ptr = warp_uniform(grad_query_ptr); + grad_key_ptr = warp_uniform(grad_key_ptr); + grad_value_ptr = warp_uniform(grad_value_ptr); + + return true; + } + + __host__ dim3 getBlocksGrid() const { + return dim3(num_splits_key, num_heads, num_batches); + } + + __host__ dim3 getThreadsGrid() const { + return dim3(kWarpSize * kNumWarpsPerBlock, 1, 1); + } + + CUTLASS_HOST_DEVICE int64_t workspace_elements_gk() const { + if constexpr (!kNeedsAccumGradK) { + return 0; + } + return static_cast(num_splits_key) * + static_cast(kBlockSizeJ) * + static_cast(align_up(head_dim, (int32_t)kBlockSizeI)); + } + + CUTLASS_HOST_DEVICE int64_t workspace_elements_gv() const { + if constexpr (!kNeedsAccumGradV) { + return 0; + } + return static_cast(num_splits_key) * + static_cast(kBlockSizeJ) * + static_cast(align_up(head_dim_value, (int32_t)kBlockSizeI)); + } + + CUTLASS_HOST_DEVICE int64_t workspace_elements_gq() const { + if constexpr (!kNeedsAccumGradQ) { + return 0; + } + auto num_blocks = + static_cast(ceil_div(num_queries, kBlockSizeI)); + auto num_cols = static_cast( + ceil_div(head_dim, MatmulGradQ::ThreadblockShape::kN)); + return num_blocks * num_cols * + static_cast(sizeof(GradQTempStorage)) / + static_cast(sizeof(output_accum_t)); + } + CUTLASS_HOST_DEVICE int64_t workspace_strideBH() const { + // Aligned on 128bits + return align_up( + workspace_elements_gk() + workspace_elements_gv() + + workspace_elements_gq(), + int64_t(4)); + } + CUTLASS_HOST_DEVICE int64_t workspace_size() const { + // Returns size of buffer we need to run this kernel + return static_cast(num_batches) * + static_cast(num_heads) * workspace_strideBH() * + static_cast(sizeof(float)); + } + CUTLASS_HOST_DEVICE bool should_zero_workspace() const { + return num_splits_key > 1; + } + }; + + struct SharedStoragePrologue { + struct { + cutlass::Array di; // (do_i * o_i).sum(-1) + cutlass::Array lse_i; // logsumexp + typename MatmulQK::Mma::SharedStorageA mm_qk_k; + } persistent; + union { + struct { + // part1 - after Q.K / dV / dO.V + union { + // 4. store Pij. it is needed: + // - in dVj += (Pij.T * Zij) @ dOi + // - in dSij = Pij * (dPij - Di) + // 6. dVj += (Pij.T * Zij) @ dOi + // 10. write to fragment + typename MatmulQK::AccumulatorSharedStorage attn_shared_storage; + }; + + union { + // 2. prologue for dVj + // 6. workspace for dVj += (Pij.T * Zij) @ dOi + typename MatmulGradV::Mma::SharedStorage mm_gradV; + // 7. dVj epilogue + typename MatmulGradV::DefaultEpilogue::SharedStorage gradV_epilogue; + }; + + // 3. prologue for dPij_dropped + // 8. used in dPij_dropped = dOi @ Vj.T + typename MatmulDOIVJ::Mma::SharedStorage mm_doivj; + } part1; + + struct { + // part2 - dQ + union { + typename MatmulQK::AccumulatorSharedStorage + tmpT_shared_storage; // (from part1) + typename MatmulDOIVJ::AccumulatorSharedStorage tmp_shared_storage; + }; + typename MatmulGradK::Mma::SharedStorage mm_gradK; // (preload) + typename MatmulGradQ::Mma::SharedStorage mm_gradQ; // (preload) + union { + // store dB = dSij to global memory + typename MatmulDOIVJ::BiasGradEpilogue::SharedStorage gradB_epilogue; + typename MatmulGradQ::DefaultEpilogue::SharedStorage gradQ_epilogue; + }; + + } part2; + + struct { + // part3 - after last iteration on dQ's epilogue / dK + union { + typename MatmulQK::AccumulatorSharedStorage + tmpT_shared_storage; // (from part1) + typename MatmulDOIVJ::AccumulatorSharedStorage tmp_shared_storage; + }; + typename MatmulGradK::Mma::SharedStorage mm_gradK; // (preload) + typename MatmulGradQ::DefaultEpilogue::SharedStorage + gradQ_epilogue_lastIter; + + typename MatmulGradK::DefaultEpilogue::SharedStorage gradK_epilogue; + } part3; + + struct { + // part4 - after last iteration on dK's epilogue / preload next K.Q_t + typename MatmulQK::Mma::SharedStorageB mm_qk_q; + + // If we reach end of current key, dump RF->gmem with "final" epilogues + typename MatmulGradK::DefaultEpilogue::SharedStorage + gradK_epilogue_final; + typename MatmulGradV::DefaultEpilogue::SharedStorage + gradV_epilogue_final; + } part4; + }; + static void print_size() { + // Field size +#define FSZ(f) int((sizeof(((SharedStoragePrologue*)0)->f))) + + printf("Total smem: %d bytes\n", int(sizeof(SharedStoragePrologue))); + printf(" persistent: %db\n", FSZ(persistent)); + printf(" mm_qk_k: %db\n", FSZ(persistent.mm_qk_k)); + printf(" part1: %db\n", FSZ(part1)); + printf(" attn_shared_storage: %db\n", FSZ(part1.attn_shared_storage)); + printf(" zij: %db\n", FSZ(part1.zij)); + printf(" mm_gradV: %db\n", FSZ(part1.mm_gradV)); + printf(" gradV_epilogue: %db\n", FSZ(part1.gradV_epilogue)); + printf(" mm_doivj: %db\n", FSZ(part1.mm_doivj)); + printf(" part2: %db\n", FSZ(part2)); + printf(" tmpT_shared_storage: %db\n", FSZ(part2.tmpT_shared_storage)); + printf(" tmp_shared_storage: %db\n", FSZ(part2.tmp_shared_storage)); + printf(" mm_gradK: %db\n", FSZ(part2.mm_gradK)); + printf(" mm_gradQ: %db\n", FSZ(part2.mm_gradQ)); + printf(" gradB_epilogue: %db\n", FSZ(part2.gradB_epilogue)); + printf(" gradQ_epilogue: %db\n", FSZ(part2.gradQ_epilogue)); + printf(" part3: %db\n", FSZ(part3)); + printf(" tmpT_shared_storage: %db\n", FSZ(part3.tmpT_shared_storage)); + printf(" part4: %db\n", FSZ(part4)); + printf(" mm_qk_q: %db\n", FSZ(part4.mm_qk_q)); + printf( + " gradK_epilogue_final: %db\n", FSZ(part4.gradK_epilogue_final)); + printf( + " gradV_epilogue_final: %db\n", FSZ(part4.gradV_epilogue_final)); + } +// =========================================== +#define FIELD(INSIDE_STRUCT, FIELDNAME) \ + CUTLASS_DEVICE auto& FIELDNAME() { \ + return INSIDE_STRUCT.FIELDNAME; \ + } + + FIELD(persistent, di) + FIELD(persistent, mm_qk_k) + FIELD(persistent, lse_i) + FIELD(part1, attn_shared_storage) + FIELD(part1, zij) + FIELD(part1, mm_gradV) + FIELD(part1, gradV_epilogue) + FIELD(part1, mm_doivj) + FIELD(part2, mm_gradK) + FIELD(part2, mm_gradQ) + FIELD(part2, gradB_epilogue) + FIELD(part2, gradQ_epilogue) + FIELD(part2, tmp_shared_storage) + FIELD(part3, tmpT_shared_storage) + FIELD(part3, gradQ_epilogue_lastIter) + FIELD(part3, gradK_epilogue) + FIELD(part4, mm_qk_q) + FIELD(part4, gradK_epilogue_final) + FIELD(part4, gradV_epilogue_final) + }; + + struct SharedStorageNoPrologue { + struct { + cutlass::Array di; // (do_i * o_i).sum(-1) + } persistent; + union { + struct { + // part1 - Q.K matmul + typename MatmulQK::Mma::SharedStorageA mm_qk_k; + typename MatmulQK::Mma::SharedStorageB mm_qk_q; + cutlass::Array lse_i; // logsumexp + } part1; + + struct { + // part2 - compute gradV + union { + // 2. store Pij to shared memory. it is needed: + // - in this step, where it is used in dVj += (Pij.T * Zij) @ dOi + // - in next step where it is used in dSij = Pij * (dPij - Di) + typename MatmulQK::AccumulatorSharedStorage attn_shared_storage; + }; + + union { + typename MatmulGradV::Mma::SharedStorage mm_gradV; + typename MatmulGradV::DefaultEpilogue::SharedStorage gradV_epilogue; + }; + } part2; + + struct { + // part3 - DO.V matmul + union { + // first compute dPij = (dOi @ Vj.T) * Zij + // and dSij = Pij * (dPij - Di) + struct { + // (from part2) - Pij for computing dSij = Pij * (dPij - Di) + typename MatmulQK::AccumulatorSharedStorage attn_shared_storage; + // matmul to compute dOiVj + typename MatmulDOIVJ::Mma::SharedStorage mm_doivj; + }; + // then store dB = dSij to global memory + typename MatmulDOIVJ::BiasGradEpilogue::SharedStorage gradB_epilogue; + }; + } part3; + + struct { + // part4 - compute gradQ + typename MatmulQK::AccumulatorSharedStorage + tmpT_shared_storage; // (from part2) + typename MatmulDOIVJ::AccumulatorSharedStorage tmp_shared_storage; + union { + typename MatmulGradQ::Mma::SharedStorage mm_gradQ; + typename MatmulGradQ::DefaultEpilogue::SharedStorage gradQ_epilogue; + typename MatmulGradQ::DefaultEpilogue::SharedStorage + gradQ_epilogue_lastIter; + }; + } part4; + + struct { + // part5 - compute gradK + typename MatmulQK::AccumulatorSharedStorage + tmpT_shared_storage; // (from part2) + typename MatmulDOIVJ::AccumulatorSharedStorage tmp_shared_storage; + union { + typename MatmulGradK::Mma::SharedStorage mm_gradK; + typename MatmulGradK::DefaultEpilogue::SharedStorage gradK_epilogue; + }; + } part5; + + struct { + // part6 - store RF accumulated into gmem + typename MatmulGradK::DefaultEpilogue::SharedStorage + gradK_epilogue_final; + typename MatmulGradV::DefaultEpilogue::SharedStorage + gradV_epilogue_final; + } part6; + }; + static void print_size() { +#define FIELD_SIZEOF(f) int((sizeof(((SharedStorageNoPrologue*)0)->f))) + printf("Total smem: %d bytes\n", int(sizeof(SharedStorageNoPrologue))); + printf(" persistent: %db\n", FIELD_SIZEOF(persistent)); + printf(" part1: %db\n", FIELD_SIZEOF(part1)); + printf(" part2: %db\n", FIELD_SIZEOF(part2)); + printf(" part3: %db\n", FIELD_SIZEOF(part3)); + printf(" part4: %db\n", FIELD_SIZEOF(part4)); + printf(" part5: %db\n", FIELD_SIZEOF(part5)); + printf(" part6: %db\n", FIELD_SIZEOF(part6)); + } +// =========================================== +#define FIELD(INSIDE_STRUCT, FIELDNAME) \ + CUTLASS_DEVICE auto& FIELDNAME() { \ + return INSIDE_STRUCT.FIELDNAME; \ + } + + FIELD(persistent, di) + FIELD(part1, lse_i) + FIELD(part1, mm_qk_k) + FIELD(part1, mm_qk_q) + FIELD(part2, attn_shared_storage) + FIELD(part2, zij) + FIELD(part2, mm_gradV) + FIELD(part2, gradV_epilogue) + FIELD(part3, mm_doivj) + FIELD(part3, gradB_epilogue) + FIELD(part4, tmpT_shared_storage) + FIELD(part4, tmp_shared_storage) + FIELD(part4, mm_gradQ) + FIELD(part4, gradQ_epilogue) + FIELD(part4, gradQ_epilogue_lastIter) + FIELD(part5, mm_gradK) + FIELD(part5, gradK_epilogue) + FIELD(part6, gradK_epilogue_final) + FIELD(part6, gradV_epilogue_final) + }; + + using SharedStorage = typename cutlass::platform::conditional< + kPreload, + SharedStoragePrologue, + SharedStorageNoPrologue>::type; + + struct OutputFragments { + typename MatmulGradV::Mma::FragmentC gradV; + typename MatmulGradK::Mma::FragmentC gradK; + + CUTLASS_DEVICE void clear() { + gradV.clear(); + gradK.clear(); + } + }; + + static bool __host__ check_supported(Params const& p) { + CHECK_ALIGNED_PTR(p.query_ptr, kMinimumAlignment); + CHECK_ALIGNED_PTR(p.key_ptr, kMinimumAlignment); + CHECK_ALIGNED_PTR(p.value_ptr, kMinimumAlignment); + CHECK_ALIGNED_PTR(p.output_ptr, kMinimumAlignment); + CHECK_ALIGNED_PTR(p.grad_output_ptr, kMinimumAlignment); + NATTEN_CHECK( + p.num_heads <= 1 || p.q_strideH % kMinimumAlignment == 0, + "query is not correctly aligned (strideH)"); + NATTEN_CHECK( + p.num_heads <= 1 || p.k_strideH % kMinimumAlignment == 0, + "key is not correctly aligned (strideH)"); + NATTEN_CHECK( + p.num_heads <= 1 || p.v_strideH % kMinimumAlignment == 0, + "value is not correctly aligned (strideH)"); + NATTEN_CHECK( + p.num_batches <= 1 || p.q_strideB % kMinimumAlignment == 0, + "query is not correctly aligned (strideB)"); + NATTEN_CHECK( + p.num_batches <= 1 || p.k_strideB % kMinimumAlignment == 0, + "key is not correctly aligned (strideB)"); + NATTEN_CHECK( + p.num_batches <= 1 || p.v_strideB % kMinimumAlignment == 0, + "value is not correctly aligned (strideB)"); + NATTEN_CHECK( + p.q_strideM % kMinimumAlignment == 0, + "query is not correctly aligned (strideM)"); + NATTEN_CHECK( + p.k_strideM % kMinimumAlignment == 0, + "key is not correctly aligned (strideM)"); + NATTEN_CHECK( + p.v_strideM % kMinimumAlignment == 0, + "value is not correctly aligned (strideM)"); + NATTEN_CHECK(p.head_dim > 0, "Invalid value for `head_dim`"); + NATTEN_CHECK(p.head_dim_value > 0, "Invalid value for `head_dim_value`"); + NATTEN_CHECK(p.num_queries > 0, "Invalid value for `num_queries`"); + NATTEN_CHECK(p.num_keys > 0, "Invalid value for `num_keys`"); + NATTEN_CHECK(p.num_heads > 0, "Invalid value for `num_heads`"); + NATTEN_CHECK(p.num_batches > 0, "Invalid value for `num_batches`"); + NATTEN_CHECK(p.head_dim <= kMaxK, "kMaxK: Expected `head_dim < kMaxK`"); + NATTEN_CHECK( + p.custom_mask_type < NumCustomMaskTypes, + "Invalid value for `custom_mask_type`"); + NATTEN_CHECK( + p.head_dim_value <= kMaxK, "kMaxK: Expected `head_dim_value < kMaxK`"); + if (kKeysQueriesAlignedToBlockSize) { + NATTEN_CHECK( + p.ptr_cumulative_seqlen_KV == nullptr, + "This kernel does not support cu_seqlen"); + NATTEN_CHECK( + p.ptr_cumulative_seqlen_Q == nullptr, + "This kernel does not support cu_seqlen"); + NATTEN_CHECK( + p.num_queries % kBlockSizeI == 0, + "kKeysQueriesAlignedToBlockSize condition not respected"); + NATTEN_CHECK( + p.num_keys % kBlockSizeJ == 0, + "kKeysQueriesAlignedToBlockSize condition not respected"); + } + NATTEN_CHECK( + kEnableSplitKeys || p.num_splits_key == 1, "SplitKeys is disabled"); + NATTEN_CHECK( + p.num_splits_key > 0, "Invalid `num_splits_key` (expected >0)"); + NATTEN_CHECK( + p.num_splits_key <= cutlass::ceil_div(p.num_keys, kBlockSizeJ), + "Invalid `num_splits_key` (" + << p.num_splits_key + << ") - too large for `num_keys` = " << p.num_keys); + return true; + } + + static CUTLASS_DEVICE void attention_kernel(Params p) { + extern __shared__ char smem_buffer[]; + SharedStorage& shared_storage = *((SharedStorage*)smem_buffer); + + uint16_t thread_id = threadIdx.x; + uint8_t warp_id = warp_uniform(thread_id / 32); + uint8_t lane_id = thread_id % 32; + + int32_t key_start = p.split_key_device() * kBlockSizeJ; + if (key_start >= p.num_keys) { + return; + } + if (kPrologueQK) { + int32_t query_start = getQueryStart(p, key_start); + prologueQkNextIteration( + shared_storage, p, query_start, key_start, warp_id, lane_id); + } + + // Computes (dO*out).sum(-1) and writes it to `p.delta_ptr` + if constexpr (kKernelComputesDelta) { + constexpr int kOptimalElements = + 128 / cutlass::sizeof_bits::value; + if (p.head_dim_value % kOptimalElements == 0) { + for (int query_start = 0; query_start < p.num_queries; + query_start += kBlockSizeI) { + computeDelta(p, query_start, warp_id, lane_id); + } + } else { + for (int query_start = 0; query_start < p.num_queries; + query_start += kBlockSizeI) { + computeDelta<1>(p, query_start, warp_id, lane_id); + } + } + __syncthreads(); + } + + OutputFragments output_frags; + + CUTLASS_PRAGMA_UNROLL + for (; key_start < p.num_keys; + key_start += p.num_splits_key_device() * kBlockSizeJ) { + output_frags.clear(); + + int32_t next_key = key_start; + int32_t query_start = getQueryStart(p, key_start); + while (next_key == key_start && query_start < p.num_queries) { + // This line here + // vvvvvvvvvvvvvv + warp_id = warp_uniform(warp_id); + // ^^^^^^^^^^^^^^ + // ... makes everything use less RF and be 10% faster. Why? + // I don't know. My theory is that it forces `nvcc` to + // re-compute indices, offsets etc... and not keep them + // from the previous iteration, which prevents MASSIVE + // register spilling. + + processBlockIJ( + shared_storage, + output_frags, + p, + query_start, + key_start, + warp_id, + lane_id); + + int32_t next_query; + incrIteration(p, query_start, key_start, next_query, next_key); + query_start = next_query; + } + if (kOutputInRF) { + writeFragsToGmem( + shared_storage, output_frags, p, key_start, warp_id, lane_id); + } else if (getQueryStart(p, key_start) >= p.num_queries) { + zfillGradKV( + p, key_start, warp_id, lane_id); + } + __syncthreads(); + } + } + + template + static CUTLASS_DEVICE void zfillGradKV( + Params const& p, + int32_t key_start, + uint8_t warp_id, + uint8_t lane_id) { + constexpr int kThreadsPerKey = 8; + constexpr int kParallelKeys = kNumThreads / kThreadsPerKey; + static_assert(kBlockSizeJ % kParallelKeys == 0, ""); + // This function is not really optimized, but should rarely be used + // It's only used when some keys are "useless" and don't attend to + // any query, due to causal masking + + int thread_id = 32 * warp_id + lane_id; + int k_shift = lane_id % kThreadsPerKey; + + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < kBlockSizeJ; j += kParallelKeys) { + int key = key_start + j + (thread_id / kThreadsPerKey); + if (!skipBoundsChecks && key >= p.num_keys) { + continue; + } + auto gv_ptr = p.grad_value_ptr + key * p.gV_strideM(); + auto gk_ptr = p.grad_key_ptr + key * p.gK_strideM(); + + for (int k = k_shift; k < p.head_dim_value; k += kThreadsPerKey) { + gv_ptr[k] = scalar_t(0); + } + for (int k = k_shift; k < p.head_dim; k += kThreadsPerKey) { + gk_ptr[k] = scalar_t(0); + } + } + } + + template + static CUTLASS_DEVICE void processBlockIJ( + SharedStorage& shared_storage, + OutputFragments& output_frags, + Params& p, + int32_t query_start, + int32_t key_start, + uint8_t warp_id, + uint8_t lane_id) { + cutlass::MatrixCoord no_offset{0, 0}; + accum_t scale = p.scale; + int16_t thread_id = 32 * warp_id + lane_id; + + auto rematerializeThreadIds = [&]() { + // Prevents `nvcc` from keeping values deduced from + // `thread_id`, `warp_id`, ... in RF - to reduce register pressure + warp_id = warp_uniform(thread_id / 32); + lane_id = thread_id % 32; + thread_id = 32 * warp_id + lane_id; + }; + + bool isFirstQuery = (query_start == getQueryStart(p, key_start)); + int32_t next_query, next_key; + incrIteration(p, query_start, key_start, next_query, next_key); + bool isLastQuery = next_key != key_start; + // PRINT_T0("[b:(%d, %d, %d)] q=%d k=%d FIRST=%d LAST=%d next: (%d, %d)", + // int(blockIdx.x), int(blockIdx.y), int(blockIdx.z), + // query_start, key_start, int(isFirstQuery), int(isLastQuery), + // int(next_query), int(next_key) + // ); + + // NOTE(alih): we fetch lse_i here as well as d_i. + accum_t lsei_rf = accum_t(0); + accum_t di_rf = accum_t(0); + if (thread_id < kBlockSizeI) { + if (query_start + thread_id < p.num_queries) { + auto query_offset = (query_start + thread_id) * p.num_heads; + di_rf = p.delta_ptr[query_offset]; + lsei_rf = p.logsumexp_ptr[query_offset]; + } + shared_storage.lse_i()[thread_id] = lsei_rf; + shared_storage.di()[thread_id] = di_rf; + } + + int32_t num_queries_in_block = skipBoundsChecks + ? MatmulQK::Mma::Shape::kN + : warp_uniform(cutlass::fast_min( + (int32_t)MatmulQK::Mma::Shape::kN, p.num_queries - query_start)); + int32_t num_keys_in_block = skipBoundsChecks + ? MatmulQK::Mma::Shape::kM + : warp_uniform(cutlass::fast_min( + (int32_t)MatmulQK::Mma::Shape::kM, p.num_keys - key_start)); + + auto prologueGradV = [&](int col) { + typename MatmulGradV::Mma::IteratorB iterator_dO( + {int32_t(p.gO_strideM)}, + p.grad_output_ptr + query_start * p.gO_strideM + col, + {num_queries_in_block, p.head_dim_value - col}, + thread_id, + no_offset); + MatmulGradV::Mma::prologue( + shared_storage.mm_gradV(), + iterator_dO, + thread_id, + num_queries_in_block); + }; + auto prologueGradQ = [&](int col) { + typename MatmulGradQ::Mma::IteratorB iterator_K( + {int32_t(p.k_strideM)}, + p.key_ptr + key_start * p.k_strideM + col, + {num_keys_in_block, p.head_dim - col}, + thread_id, + no_offset); + MatmulGradQ::Mma::prologue( + shared_storage.mm_gradQ(), iterator_K, thread_id, num_keys_in_block); + }; + auto prologueGradK = [&](int col) { + typename MatmulGradK::Mma::IteratorB iterator_Q( + {int32_t(p.q_strideM)}, + p.query_ptr + query_start * p.q_strideM + col, + {num_queries_in_block, p.head_dim - col}, + thread_id, + no_offset); + MatmulGradK::Mma::prologue( + shared_storage.mm_gradK(), + iterator_Q, + thread_id, + num_queries_in_block); + }; + auto prologueDOV = [&]() { + typename MatmulDOIVJ::Mma::IteratorA iterator_A( + {int32_t(p.gO_strideM)}, + p.grad_output_ptr + query_start * p.gO_strideM, + {num_queries_in_block, p.head_dim_value}, + thread_id, + no_offset); + typename MatmulDOIVJ::Mma::IteratorB iterator_B( + {int32_t(p.v_strideM)}, + p.value_ptr + key_start * p.v_strideM, + {p.head_dim_value, num_keys_in_block}, + thread_id, + no_offset); + MatmulDOIVJ::Mma::prologue( + shared_storage.mm_doivj(), + iterator_A, + iterator_B, + thread_id, + p.head_dim_value); + }; + + ///////////////////////////////////////////////////////////////////////////////////////////////// + // MatmulQK + ///////////////////////////////////////////////////////////////////////////////////////////////// + { + using Mma = typename MatmulQK::Mma; + + cutlass::gemm::GemmCoord problem_size( + num_keys_in_block, + num_queries_in_block, + p.head_dim // k + ); + + // k_j + typename Mma::IteratorA iterator_A( + {int32_t(p.k_strideM)}, + p.key_ptr + key_start * p.k_strideM, + {problem_size.m(), problem_size.k()}, + thread_id, + no_offset); + + // q_i.transpose(-2, -1) + typename Mma::IteratorB iterator_B( + {int32_t(p.q_strideM)}, + p.query_ptr + query_start * p.q_strideM, + {problem_size.k(), problem_size.n()}, + thread_id, + no_offset); + + Mma mma( + shared_storage.mm_qk_k(), + shared_storage.mm_qk_q(), + thread_id, + warp_id, + lane_id); + + typename Mma::FragmentC accum; + + accum.clear(); + + auto gemm_k_iterations = + (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Compute threadblock-scoped matrix multiply-add + mma.set_prologue_done(kPrologueQK); + mma.set_zero_outside_bounds(!skipBoundsChecks); + mma(gemm_k_iterations, accum, iterator_A, iterator_B, accum); + accum = cutlass::multiplies()(scale, accum); + + // Epilogue: add LSE + exp and store that to our shared memory buffer + // shmem <- (matmul_result - + // logsumexp[i_start:i_end].unsqueeze(1)).exp() + int warp_idx_mn_0 = + warp_id % (Mma::Base::WarpCount::kM * Mma::Base::WarpCount::kN); + auto output_tile_coords = cutlass::MatrixCoord{ + warp_idx_mn_0 % Mma::Base::WarpCount::kM, + warp_idx_mn_0 / Mma::Base::WarpCount::kM}; + + auto lane_offset = MatmulQK::AccumLambdaIterator::get_lane_offset( + lane_id, warp_id, output_tile_coords); + + MatmulQK::AccumLambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) {}, + [&](int accum_m, int accum_n, int idx) { + // NOTE(alih): Checking whether the current q coordinate is out of + // bounds is something I had to do for SM70 and SM50 kernels in FNA + // as well. Unsure why / and why xFormers doesn't (didn't) run into + // issues with this. + if (accum_m + key_start >= p.num_keys +#if __CUDA_ARCH__ < 800 + || accum_n + query_start >= p.num_queries +#endif + ) { + accum[idx] = + -cutlass::platform::numeric_limits::infinity(); + } else { + accum[idx] -= shared_storage.lse_i()[accum_n]; + } + }, + [&](int accum_m) {}); + + // Apply mask + if (p.custom_mask_type == CausalFromTopLeft || + p.custom_mask_type == CausalFromBottomRight) { + auto lane_offset = MatmulQK::AccumLambdaIterator::get_lane_offset( + lane_id, warp_id, output_tile_coords); + int shift = query_start - key_start; + if (p.custom_mask_type == CausalFromBottomRight) { + shift += p.num_keys - p.num_queries; + } + // current_key = key_start + accum_m + // current_query = query_start + accum_n + // mask if: `current_key > current_query` + MatmulQK::AccumLambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) {}, + [&](int accum_m, int accum_n, int idx) { + if (accum_m > accum_n + shift) { + accum[idx] = + -cutlass::platform::numeric_limits::infinity(); + } + }, + [&](int accum_m) {}); + } + + __syncthreads(); + if (kPrologueGV) { + prologueGradV(0); + } + if (kPrologueDOV) { + prologueDOV(); + } + + // MatmulQK::B2bGemm::accumApplyLSEToSmem( + MatmulQK::B2bGemm::accumApplyExpToSmem( + shared_storage.attn_shared_storage(), + accum, + // p.logsumexp_ptr + query_start * p.num_heads, + // problem_size.n(), + thread_id, + warp_id, + lane_id, + output_tile_coords); + + __syncthreads(); + } + rematerializeThreadIds(); + + ///////////////////////////////////////////////////////////////////////////////////////////////// + // GradV matmul + // + // grad_v[j_start:j_end] += attn_T @ do_i + ///////////////////////////////////////////////////////////////////////////////////////////////// + constexpr bool kSingleIterationGradV = + kMaxK <= MatmulGradV::ThreadblockShape::kN; + for (int col = 0; col < (kSingleIterationGradV ? 1 : p.head_dim_value); + col += MatmulGradV::ThreadblockShape::kN) { + using Mma = typename MatmulGradV::Mma; + using AccumTileGmem = typename MatmulGradQ::AccumTileGmem; + + cutlass::gemm::GemmCoord problem_size( + num_keys_in_block, p.head_dim_value - col, num_queries_in_block); + auto createEpilogueIter = [&]() { + return typename MatmulGradV::OutputTileIterator( + typename MatmulGradV::OutputTileIterator::Params{p.gV_strideM()}, + p.grad_value_ptr + key_start * p.gV_strideM() + col, + {num_keys_in_block, p.head_dim_value - col}, + thread_id); + }; + typename Mma::IteratorB iterator_B( + {int32_t(p.gO_strideM)}, + p.grad_output_ptr + query_start * p.gO_strideM + col, + {num_queries_in_block, p.head_dim_value - col}, + thread_id, + no_offset); + + // dVj += Pij.T @ dOi + Mma mma( + // operand A: Pij.T + shared_storage.attn_shared_storage().accum_ref(), + //// operand A_scale Zij.T: + //// if we're using dropout, operand A is Pij_dropped.T = Pij.T * + /// Zij.T / which is computed on the fly as fragments of Pij.T are + /// loaded in + // shared_storage.zij().accum_ref(), + // operand B: dOi - which was loaded into shared memory previously + // when we computed dVj + shared_storage.mm_gradV().operand_B_ref(), + thread_id, + warp_id, + lane_id); + + int storage_id = col / MatmulGradV::ThreadblockShape::kN; + AccumTileGmem gmem_tile{ + p.workspace_gv + storage_id * AccumTileGmem::kElementsStored}; + if (!kOutputInRF) { + if (isFirstQuery || !kNeedsAccumGradV) { + output_frags.gradV.clear(); + } else { + gmem_tile.load(output_frags.gradV, thread_id); + } + } + mma.set_prologue_done(kPrologueGV); + + auto gemm_k_iterations = + (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Compute threadblock-scoped matrix multiply-add + __syncthreads(); + + mma(gemm_k_iterations, + output_frags.gradV, + iterator_B, + output_frags.gradV); + __syncthreads(); + if (kPrologueGV && !kSingleIterationGradV && + col + MatmulGradV::ThreadblockShape::kN < p.head_dim_value) { + prologueGradV(col + MatmulGradV::ThreadblockShape::kN); + } + + if (!kOutputInRF) { + if (kNeedsAccumGradV && !isLastQuery) { + gmem_tile.store(output_frags.gradV, thread_id); + } else { + accumulateInGmem( + shared_storage.gradV_epilogue(), + output_frags.gradV, + createEpilogueIter(), + isFirstQuery || kNeedsAccumGradV, + warp_id, + lane_id); + } + } + } + __syncthreads(); + + ///////////////////////////////////////////////////////////////////////////////////////////////// + // MatmulDOIVJ + ///////////////////////////////////////////////////////////////////////////////////////////////// + { + using Mma = typename MatmulDOIVJ::Mma; + // do_i + typename Mma::IteratorA iterator_A( + {int32_t(p.gO_strideM)}, + p.grad_output_ptr + query_start * p.gO_strideM, + {num_queries_in_block, p.head_dim_value}, + thread_id, + no_offset); + + // v_j.transpose(-2, -1) + typename Mma::IteratorB iterator_B( + {int32_t(p.v_strideM)}, + p.value_ptr + key_start * p.v_strideM, + {p.head_dim_value, num_keys_in_block}, + thread_id, + no_offset); + + Mma mma(shared_storage.mm_doivj(), thread_id, warp_id, lane_id); + mma.set_prologue_done(kPrologueDOV); + mma.set_zero_outside_bounds(!skipBoundsChecks); + + typename Mma::FragmentC accum; + + accum.clear(); + + auto gemm_k_iterations = + (p.head_dim_value + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Compute threadblock-scoped matrix multiply-add + mma(gemm_k_iterations, accum, iterator_A, iterator_B, accum); + __syncthreads(); + if (kPrologueGQ) { + prologueGradQ(0); + } + if (kPrologueGK) { + prologueGradK(0); + } + + int warp_idx_mn_0 = + warp_id % (Mma::Base::WarpCount::kM * Mma::Base::WarpCount::kN); + auto output_tile_coords = cutlass::MatrixCoord{ + warp_idx_mn_0 % Mma::Base::WarpCount::kM, + warp_idx_mn_0 / Mma::Base::WarpCount::kM}; + // TODO: This must be terribly inefficient. There must be a better way + // tmp [RF] <- (accum [RF] - Di [smem] ) * attn_T.T [smem] + // attn_shared_storage [smem] <- tmp.T + // tmp_shared_storage [smem] <- tmp + { + using LambdaIterator = typename MatmulDOIVJ::AccumLambdaIterator; + auto lane_offset = LambdaIterator::get_lane_offset( + lane_id, warp_id, output_tile_coords); + + auto attn_T = shared_storage.attn_shared_storage().accum_ref(); + + accum_t current_di; + // dSij = (dPij - Di) * Pij + LambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) { current_di = shared_storage.di()[accum_m]; }, + [&](int accum_m, int accum_n, int idx) { + // TODO: Otherwise we can get nans as we + // might have infs here (only seen on f16 tho) + if (skipBoundsChecks || + (accum_m < num_queries_in_block && + accum_n < num_keys_in_block)) { + accum_t attn = attn_T.at({accum_n, accum_m}); + accum[idx] = (accum[idx] - current_di) * attn; + } else { + accum[idx] = 0; + } + }, + [&](int accum_m) { + + }); + + accum = accum * scale; + +#if 0 + PRINT_B0_T0("(doivj - di) * attn * scale"); + print_warp_accum(accum, lane_offset, 4, 4); +#endif + + __syncthreads(); + if (!MatmulGradK::DefaultMmaFromSmem::kIsTransposedA) { + auto tmpT = shared_storage.tmpT_shared_storage().accum_ref(); + // attn <- attn_T.T + LambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) {}, + [&](int accum_m, int accum_n, int idx) { + tmpT.at({accum_n, accum_m}) = scalar_t(accum[idx]); + }, + [&](int accum_m) {}); + } + } + + MatmulDOIVJ::B2bGemm::accumToSmem( + shared_storage.tmp_shared_storage(), + accum, + lane_id, + output_tile_coords); + __syncthreads(); + } + // Force `nvcc` to recompute values that depend on the variables just below + // to use less RF and prevent some spilling + p.head_dim = warp_uniform(p.head_dim); + p.k_strideM = warp_uniform(p.k_strideM); + rematerializeThreadIds(); + + ///////////////////////////////////////////////////////////////////////////////////////////////// + // GradQ matmul + // + // grad_q[i_start:i_end] += tmp @ k_j + ///////////////////////////////////////////////////////////////////////////////////////////////// + // Skip the loop & associated branches if we know at compile time the number + // of iterations + constexpr bool kSingleIterationGradQ = + kMaxK <= MatmulGradQ::ThreadblockShape::kN; + for (int col = 0; col < (kSingleIterationGradQ ? 1 : p.head_dim); + col += MatmulGradQ::ThreadblockShape::kN) { + using Mma = typename MatmulGradQ::Mma; + using AccumTileGmem = typename MatmulGradQ::AccumTileGmem; + + cutlass::gemm::GemmCoord problem_size( + num_queries_in_block, + false ? MatmulGradQ::ThreadblockShape::kN : p.head_dim - col, + num_keys_in_block); + + // k_j + typename Mma::IteratorB iterator_B( + {int32_t(p.k_strideM)}, + p.key_ptr + key_start * p.k_strideM + col, + {problem_size.k(), problem_size.n()}, + thread_id, + no_offset); + + auto a = shared_storage.tmp_shared_storage().accum_ref(); + Mma mma( + // operand A: dSij + shared_storage.tmp_shared_storage().accum_ref(), + // operand B: Kj + shared_storage.mm_gradQ().operand_B_ref(), + thread_id, + warp_id, + lane_id); + + typename Mma::FragmentC accum; + + int col_id = col / MatmulGradQ::ThreadblockShape::kN; + int num_cols = kSingleIterationGradQ + ? 1 + : ceil_div(p.head_dim, MatmulGradQ::ThreadblockShape::kN); + int storage_id = (col_id + query_start / kBlockSizeI * num_cols); + + if (p.num_splits_key_device() > 1) { + AtomicLock::acquire( + &p.workspace_gq[storage_id].lock, + p.split_key_device() + 1, + thread_id); + // Make sure we can see other block's output + __threadfence(); + } + + AccumTileGmem gmem_tile{&p.workspace_gq[storage_id].buffer[0]}; + if (!kNeedsAccumGradQ || + (p.num_splits_key_device() == 1 && key_start == 0)) { + // if we know we are the first to access it, we know it's only zeros. + // Avoids a load from gmem (and gmem init as well) + accum.clear(); + } else { + gmem_tile.load(accum, thread_id); + } + + auto gemm_k_iterations = + (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Compute threadblock-scoped matrix multiply-add + __syncthreads(); + mma.set_prologue_done(kPrologueGQ); + mma(gemm_k_iterations, accum, iterator_B, accum); + __syncthreads(); + bool isLastColumn = kSingleIterationGradQ || + (col + MatmulGradQ::ThreadblockShape::kN >= p.head_dim); + if (kPrologueGQ && !isLastColumn) { + prologueGradQ(col + MatmulGradQ::ThreadblockShape::kN); + } + + bool isLast = [&]() { + int32_t next_key = key_start + p.num_splits_key_device() * kBlockSizeJ; + if (p.num_keys <= next_key) { + return true; + } + if (query_start < getSmallestQueryForKey(p, next_key)) { + return true; + } + return false; + }(); + // Output results + if (p.num_splits_key_device() > 1) { + int32_t numAddsSoFar = -1; + if (isLast && thread_id == 0) { + numAddsSoFar = atomicAdd(&p.workspace_gq[storage_id].counter, 1) + + 1; // `atomicAdd` returns the old value + } + isLast = __syncthreads_or( + numAddsSoFar == getNumParallelBlocksForQuery(p, query_start)); + assert(numAddsSoFar <= getNumParallelBlocksForQuery(p, query_start)); + } + if (kNeedsAccumGradQ && !isLast) { + gmem_tile.store(accum, thread_id); + if (p.num_splits_key_device() > 1) { + // Make sure everyone wrote before we release the lock + __threadfence(); + __syncthreads(); + AtomicLock::release(&p.workspace_gq[storage_id].lock, thread_id); + } + } else { + // NOTE: We're not releasing the lock because no one is expected + // to come after us (we're the last one to write) + typename MatmulGradQ::OutputTileIterator output_it( + typename MatmulGradQ::OutputTileIterator::Params{p.gQ_strideM()}, + p.grad_query_ptr + query_start * p.gQ_strideM() + col, + {problem_size.m(), problem_size.n()}, + thread_id); + // if `direct_store` is True, we store to gmem (`*gmem = accum`) + // otherwise, we accumulate in gmem (`*gmem = *gmem + accum`) + // If we know ahead of time when we will write for the first time + // we can: + // (1) Avoid an additional memory read + // (2) Avoid the cost of initializing memory to 0 + bool direct_store = kNeedsAccumGradQ || key_start == 0 || + (p.num_splits_key_device() > 1); + accumulateInGmem( + isLastColumn ? shared_storage.gradQ_epilogue_lastIter() + : shared_storage.gradQ_epilogue(), + accum, + output_it, + direct_store, + warp_id, + lane_id); + } + } + ///////////////////////////////////////////////////////////////////////////////////////////////// + // GradK matmul + // + // grad_k[i_start:i_end] += tmp.transpose(-2, -1) @ q_i + ///////////////////////////////////////////////////////////////////////////////////////////////// + rematerializeThreadIds(); + + constexpr bool kSingleIterationGradK = + kMaxK <= MatmulGradK::ThreadblockShape::kN; + for (int col = 0; col < (kSingleIterationGradK ? 1 : p.head_dim); + col += MatmulGradK::ThreadblockShape::kN) { + using Mma = typename MatmulGradK::Mma; + using AccumTileGmem = typename MatmulGradQ::AccumTileGmem; + + cutlass::gemm::GemmCoord problem_size( + num_keys_in_block, + false ? MatmulGradK::ThreadblockShape::kN : p.head_dim - col, + num_queries_in_block); + auto createEpilogueIter = [&]() { + return typename MatmulGradK::OutputTileIterator( + typename MatmulGradK::OutputTileIterator::Params{p.gK_strideM()}, + p.grad_key_ptr + key_start * p.gK_strideM() + col, + {num_keys_in_block, + false ? MatmulGradK::ThreadblockShape::kN : p.head_dim - col}, + thread_id); + }; + + // q_i + typename Mma::IteratorB iterator_B( + {int32_t(p.q_strideM)}, + p.query_ptr + query_start * p.q_strideM + col, + {problem_size.k(), problem_size.n()}, + thread_id, + no_offset); + + auto getTmp = [&](int) { return &shared_storage.tmp_shared_storage(); }; + auto getTmpT = [&](int) { return &shared_storage.tmpT_shared_storage(); }; + // this is basically: + // opA = kIsTransposedA ? getTmp() : getTmpT(); + bool constexpr kIsTransposedA = + MatmulGradK::DefaultMmaFromSmem::kIsTransposedA; + auto& opA = *call_conditional< + kIsTransposedA, + decltype(getTmp), + decltype(getTmpT)>::apply(getTmp, getTmpT, 0); + Mma mma( + // operand A: dSij.T + opA.accum_ref(), + // operand B: Qi + shared_storage.mm_gradK().operand_B_ref(), + thread_id, + warp_id, + lane_id); + + int storage_id = col / MatmulGradK::ThreadblockShape::kN; + AccumTileGmem gmem_tile{ + p.workspace + storage_id * AccumTileGmem::kElementsStored}; + if (!kOutputInRF) { + if (isFirstQuery || !kNeedsAccumGradK) { + output_frags.gradK.clear(); + } else { + gmem_tile.load(output_frags.gradK, thread_id); + } + } + mma.set_prologue_done(kPrologueGK); + + auto gemm_k_iterations = + (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Compute threadblock-scoped matrix multiply-add + __syncthreads(); + + mma(gemm_k_iterations, + output_frags.gradK, + iterator_B, + output_frags.gradK); + __syncthreads(); + bool isLastColumn = kSingleIterationGradK || + col + MatmulGradK::ThreadblockShape::kN >= p.head_dim; + if (kPrologueGK && !isLastColumn) { + prologueGradK(col + MatmulGradK::ThreadblockShape::kN); + } + + if (kPrologueQK && isLastColumn) { + int32_t next_query, next_key; + incrIteration(p, query_start, key_start, next_query, next_key); + DISPATCH_BOOL( + next_key != key_start, kForceReloadK, ([&]() { + prologueQkNextIteration( + shared_storage, p, next_query, next_key, warp_id, lane_id); + })); + } + + // Output results + if (!kOutputInRF) { + if (kNeedsAccumGradK && !isLastQuery) { + gmem_tile.store(output_frags.gradK, thread_id); + } else { + accumulateInGmem( + isLastColumn ? shared_storage.gradK_epilogue_final() + : shared_storage.gradK_epilogue(), + output_frags.gradK, + createEpilogueIter(), + isFirstQuery || kNeedsAccumGradK, + warp_id, + lane_id); + __syncthreads(); + } + } + } + } + + static CUTLASS_HOST_DEVICE int32_t getQueryStartShift(Params const& p) { + if (p.custom_mask_type == NoCustomMask && p.num_splits_key_device() > 1) { + return (p.split_key_device() * kBlockSizeI) % getQueryEnd(p); + } + return 0; + } + + // Iteration order logic + static CUTLASS_HOST_DEVICE int32_t + getQueryStart(Params const& p, int32_t key_start) { + return getSmallestQueryForKey(p, key_start) + getQueryStartShift(p); + }; + static CUTLASS_HOST_DEVICE int32_t getQueryEnd(Params const& p) { + return align_up(p.num_queries, kBlockSizeI); + }; + + static CUTLASS_HOST_DEVICE int32_t + getSmallestQueryForKey(Params const& p, int32_t key_start) { + if (p.custom_mask_type == NoCustomMask) { + return 0; + } + int32_t shift = p.custom_mask_type == CausalFromBottomRight + ? p.num_keys - p.num_queries + : 0; + int32_t window_size = p.num_queries + p.num_keys; + + auto last_key_for_block = + cutlass::fast_min(key_start + kBlockSizeJ, p.num_keys) - 1; + int first_query = key_start - shift; + int last_query = last_key_for_block - shift + window_size - 1; + if (last_query < 0 || first_query >= p.num_queries) { + return getQueryEnd(p); // nothing to compute in this column + } + first_query = cutlass::fast_max(0, first_query); + return (first_query / kBlockSizeI) * kBlockSizeI; + }; + + // Returns how many kernel blocks will write to a given block in `grad_query` + // This is usually equal to the number of key splits, but can be different + // for instance in the causal case, or varying seqlen + static CUTLASS_HOST_DEVICE int32_t + getNumParallelBlocksForQuery(Params const& p, int32_t query_start) { + int16_t num_key_blocks = ceil_div(p.num_keys, kBlockSizeJ); + if (p.custom_mask_type != NoCustomMask) { + int32_t shift = p.custom_mask_type == CausalFromBottomRight + ? p.num_keys - p.num_queries + : 0; + int32_t last_query_for_block = + cutlass::fast_min(query_start + kBlockSizeI, p.num_queries) - 1; + int32_t last_key_for_block = + cutlass::fast_min(last_query_for_block + shift, p.num_keys - 1); + int32_t first_key_for_block = 0; + + num_key_blocks = last_key_for_block / kBlockSizeJ + 1; + + if (last_key_for_block < 0 || first_key_for_block >= p.num_keys) { + num_key_blocks = 0; + } + } + return cutlass::fast_min(p.num_splits_key_device(), num_key_blocks); + }; + + // Returns the next block to process + static CUTLASS_HOST_DEVICE void incrIteration( + Params const& p, + int32_t query_start, + int32_t key_start, + int32_t& next_query, + int32_t& next_key) { + next_query = query_start + kBlockSizeI; + next_key = key_start; + auto query_shift = getQueryStartShift(p); + // Wrap around + if (query_shift) { + if (next_query >= p.num_queries) { + next_query = getSmallestQueryForKey(p, key_start); + return; + } else if (query_start < query_shift && query_shift <= next_query) { + // jump to next key + } else { + return; + } + } else { + if (next_query < p.num_queries) { + return; + } + // jump to next key + } + // Next key + next_key = key_start + p.num_splits_key_device() * kBlockSizeJ; + next_query = getQueryStart(p, next_key); + } + + template + static CUTLASS_DEVICE void prologueQkNextIteration( + SharedStorage& shared_storage, + Params const& p, + int32_t query_start, + int32_t key_start, + uint8_t warp_id, + uint8_t lane_id) { + if (query_start >= p.num_queries || key_start >= p.num_keys) { + return; + } + + static constexpr bool kReloadK = + kForceReloadK || !MatmulQK::Mma::kSmemContainsEntireMat; + int thread_id = 32 * warp_id + lane_id; + typename MatmulQK::Mma::IteratorA iterator_A( + {int32_t(p.k_strideM)}, + p.key_ptr + key_start * p.k_strideM, + {p.num_keys - key_start, p.head_dim}, + thread_id, + cutlass::MatrixCoord{0, 0}); + + typename MatmulQK::Mma::IteratorB iterator_B( + {int32_t(p.q_strideM)}, + p.query_ptr + query_start * p.q_strideM, + {p.head_dim, p.num_queries - query_start}, + thread_id, + cutlass::MatrixCoord{0, 0}); + + MatmulQK::Mma::prologue( + shared_storage.mm_qk_k(), + shared_storage.mm_qk_q(), + iterator_A, + iterator_B, + thread_id, + p.head_dim); + } + + template + static CUTLASS_DEVICE void writeFragsToGmem( + SharedStorage& shared_storage, + OutputFragments& output_frags, + Params const& p, + int32_t key_start, + uint8_t warp_id, + uint8_t lane_id) { + uint16_t thread_id = 32 * warp_id + lane_id; + int32_t num_keys_in_block = skipBoundsChecks + ? MatmulQK::Mma::Shape::kM + : cutlass::fast_min( + (int32_t)MatmulQK::Mma::Shape::kM, p.num_keys - key_start); + typename MatmulGradV::OutputTileIterator outputV_it( + typename MatmulGradV::OutputTileIterator::Params{p.gV_strideM()}, + p.grad_value_ptr + key_start * p.gV_strideM(), + {num_keys_in_block, p.head_dim_value}, + thread_id); + accumulateInGmem( + shared_storage.gradV_epilogue_final(), + output_frags.gradV, + outputV_it, + true, + warp_id, + lane_id); + + typename MatmulGradK::OutputTileIterator outputK_it( + typename MatmulGradK::OutputTileIterator::Params{p.gK_strideM()}, + p.grad_key_ptr + key_start * p.gK_strideM(), + {num_keys_in_block, + false ? MatmulGradK::ThreadblockShape::kN : p.head_dim}, + thread_id); + accumulateInGmem( + shared_storage.gradK_epilogue_final(), + output_frags.gradK, + outputK_it, + true, + warp_id, + lane_id); + } + + template + static CUTLASS_DEVICE void accumulateInGmem( + typename MatmulT::DefaultEpilogue::SharedStorage& epilogue_smem, + typename MatmulT::Mma::FragmentC const& accum, + typename MatmulT::OutputTileIterator output_it, + bool first, + uint8_t warp_id, + uint8_t lane_id) { + using DefaultEpilogue = typename MatmulT::DefaultEpilogue; + using DefaultOutputOp = typename MatmulT::DefaultOutputOp; + using Mma = typename MatmulT::Mma; + int thread_id = 32 * warp_id + lane_id; + DISPATCH_BOOL( + first, kIsFirst, ([&]() { + static constexpr auto ScaleType = kIsFirst + ? cutlass::epilogue::thread::ScaleType::Nothing + : cutlass::epilogue::thread::ScaleType::NoBetaScaling; + using EpilogueOutputOp = + typename cutlass::epilogue::thread::LinearCombination< + typename DefaultOutputOp::ElementOutput, + DefaultOutputOp::kCount, + typename DefaultOutputOp::ElementAccumulator, + typename DefaultOutputOp::ElementCompute, + ScaleType>; + using Epilogue = + typename cutlass::epilogue::threadblock::EpiloguePipelined< + typename DefaultEpilogue::Shape, + typename Mma::Operator, + DefaultEpilogue::kPartitionsK, + typename MatmulT::OutputTileIterator, + typename DefaultEpilogue::AccumulatorFragmentIterator, + typename DefaultEpilogue::WarpTileIterator, + typename DefaultEpilogue::SharedLoadIterator, + EpilogueOutputOp, + typename DefaultEpilogue::Padding, + DefaultEpilogue::kFragmentsPerIteration, + true // IterationsUnroll + >; + EpilogueOutputOp rescale({1, 1}); + Epilogue epilogue(epilogue_smem, thread_id, warp_id, lane_id); + epilogue(rescale, output_it, accum, output_it); + })); + } + + template + static CUTLASS_DEVICE void computeDelta( + Params const& p, + int32_t query_start, + uint8_t warp_id, + uint8_t lane_id) { + // Each thread computes one value for Delta + // Depending on warp configuration, we might have multiple + // threads of the same warp working on the same row + using AccessType = cutlass::Array; + static_assert(kNumThreads >= kBlockSizeI, ""); + static constexpr int kNumThreadsPerLine = kNumThreads / kBlockSizeI; + int16_t thread_id = 32 * warp_id + lane_id; + + int16_t laneFirstCol = kElementsPerAccess * (lane_id % kNumThreadsPerLine); + int16_t laneRow = thread_id / kNumThreadsPerLine; + bool rowPred = (query_start + laneRow) < p.num_queries; + bool pred = rowPred; + + // on windows, previous syntax __restrict__ AccessType* + // resulted in error: "restrict" is not allowed + const AccessType* __restrict__ grad_output_ptr = + reinterpret_cast( + p.grad_output_ptr + (query_start + laneRow) * p.gO_strideM + + laneFirstCol); + const AccessType* __restrict__ output_ptr = + reinterpret_cast( + p.output_ptr + (query_start + laneRow) * p.o_strideM() + + laneFirstCol); + + static constexpr int64_t kMaxIters = + kMaxK / (kElementsPerAccess * kNumThreadsPerLine); + constexpr int kPipelineStages = 2; + accum_t delta_value = accum_t(0); + using GlobalLoad = + cutlass::arch::global_load; + AccessType frag_grad_output[kPipelineStages]; + AccessType frag_output[kPipelineStages]; + + auto loadAndIncrement = [&](int ld_pos, bool is_valid) { + frag_grad_output[ld_pos].clear(); + frag_output[ld_pos].clear(); + GlobalLoad(frag_grad_output[ld_pos], grad_output_ptr, is_valid); + GlobalLoad(frag_output[ld_pos], output_ptr, is_valid); + grad_output_ptr += kNumThreadsPerLine; + output_ptr += kNumThreadsPerLine; + }; + + CUTLASS_PRAGMA_UNROLL + for (int iter = 0; iter < kPipelineStages - 1; ++iter) { + int ld_pos = iter % kPipelineStages; + pred = pred && + (laneFirstCol + iter * kElementsPerAccess * kNumThreadsPerLine) < + p.head_dim_value; + loadAndIncrement(ld_pos, pred); + } + auto columnIteration = [&](int iter) { + // Load for next iter + int ld_pos = (iter + kPipelineStages - 1) % kPipelineStages; + pred = pred && + (laneFirstCol + + (iter + kPipelineStages - 1) * kElementsPerAccess * + kNumThreadsPerLine) < p.head_dim_value; + loadAndIncrement(ld_pos, pred); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < AccessType::kElements; ++i) { + delta_value += accum_t(frag_output[iter % kPipelineStages][i]) * + accum_t(frag_grad_output[iter % kPipelineStages][i]); + } + }; + + // If we have a small lower-bound for K, we can unroll the loop + if (kMaxK <= 256) { + CUTLASS_PRAGMA_UNROLL + for (int iter = 0; iter < kMaxIters; ++iter) { + columnIteration(iter); + } + } else { + int num_iters = + ceil_div(p.head_dim_value, kElementsPerAccess * kNumThreadsPerLine) * + (kElementsPerAccess * kNumThreadsPerLine); + for (int iter = 0; iter < num_iters; ++iter) { + columnIteration(iter); + } + } + + // Reduce between workers + static_assert( + kNumThreadsPerLine == 1 || kNumThreadsPerLine == 2 || + kNumThreadsPerLine == 4, + ""); + CUTLASS_PRAGMA_UNROLL + for (int i = 1; i < kNumThreadsPerLine; i *= 2) { + delta_value = delta_value + __shfl_xor_sync(0xffffffff, delta_value, i); + } + + // Store in gmem + if (rowPred) { + p.delta_ptr[(query_start + laneRow) * p.num_heads] = delta_value; + } + } +}; + +} // namespace fmha +} // namespace cuda +} // namespace natten diff --git a/natten/csrc/include/natten/cuda/fmha/kernel_forward.h b/natten/csrc/include/natten/cuda/fmha/kernel_forward.h new file mode 100644 index 00000000..751f5158 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha/kernel_forward.h @@ -0,0 +1,1075 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ +#pragma once + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace natten { +namespace cuda { +namespace fmha { + +using namespace gemm_kernel_utils; + +namespace { +template +constexpr int getWarpsPerSmFw() { + return ( + Arch::kMinComputeCapability >= 80 && + !cutlass::platform::is_same::value + ? 16 + : 12); +} +static CUTLASS_DEVICE float atomicMaxFloat(float* addr, float value) { + // source: https://stackoverflow.com/a/51549250 + return (value >= 0) + ? __int_as_float(atomicMax((int*)addr, __float_as_int(value))) + : __uint_as_float(atomicMin((unsigned int*)addr, __float_as_uint(value))); +} + +} // namespace + +template < + // The datatype of Q/K/V + typename scalar_t_, + // Architecture we are targeting (eg `cutlass::arch::Sm80`) + typename ArchTag, + // If Q/K/V are correctly aligned in memory and we can run a fast kernel + bool isAligned_, + int kQueriesPerBlock_, + int kKeysPerBlock_, + // upperbound on `max(value.shape[-1], query.shape[-1])` + int kMaxK_ = (int)cutlass::platform::numeric_limits::max()> +struct AttentionKernel { + enum CustomMaskType { + NoCustomMask = 0, + CausalFromTopLeft = 1, + CausalFromBottomRight = 2, + NumCustomMaskTypes, + }; + using scalar_t = scalar_t_; + using accum_t = float; + using lse_scalar_t = float; + using output_t = scalar_t; + // Accumulator between 2 iterations + // Using `accum_t` improves perf on f16 at the cost of + // numerical errors + using output_accum_t = accum_t; + static constexpr int kKeysPerBlock = kKeysPerBlock_; + static constexpr int kQueriesPerBlock = kQueriesPerBlock_; + static constexpr int kMaxK = kMaxK_; + static constexpr bool kIsAligned = isAligned_; + static constexpr bool kSingleValueIteration = kMaxK <= kKeysPerBlock; + static constexpr int32_t kAlignLSE = 32; // block size of backward + static constexpr bool kIsHalf = cutlass::sizeof_bits::value == 16; + static constexpr bool kPreloadV = + ArchTag::kMinComputeCapability >= 80 && kIsHalf; + static constexpr bool kKeepOutputInRF = kSingleValueIteration; + static constexpr bool kNeedsOutputAccumulatorBuffer = !kKeepOutputInRF && + !cutlass::platform::is_same::value; + + static_assert(kQueriesPerBlock % 32 == 0, ""); + static_assert(kKeysPerBlock % 32 == 0, ""); + static constexpr int kNumWarpsPerBlock = + kQueriesPerBlock * kKeysPerBlock / (32 * 32); + static constexpr int kWarpSize = 32; + + // Launch bounds + static constexpr int kNumThreads = kWarpSize * kNumWarpsPerBlock; + static constexpr int kMinBlocksPerSm = + getWarpsPerSmFw() / kNumWarpsPerBlock; + + struct Params { + // Input tensors + const scalar_t* query_ptr = nullptr; // [num_queries, num_heads, head_dim] + const scalar_t* key_ptr = nullptr; // [num_keys, num_heads, head_dim] + const scalar_t* value_ptr = + nullptr; // [num_keys, num_heads, head_dim_value] + + const int32_t* ptr_cumulative_seqlen_Q = nullptr; + const int32_t* ptr_cumulative_seqlen_KV = nullptr; + uint32_t causal_diagonal_offset = 0; + + // Output tensors + output_t* output_ptr = nullptr; // [num_queries, num_heads, head_dim_value] + // [num_queries, num_heads, head_dim_value] + output_accum_t* output_accum_ptr = nullptr; + // [num_queries, num_heads] - can be null + lse_scalar_t* logsumexp_ptr = nullptr; + + // Scale + accum_t scale = 0.0; + + // Dimensions/strides + int32_t head_dim = 0; + int32_t head_dim_value = 0; + int32_t num_queries = 0; + int32_t num_keys = 0; + + uint8_t custom_mask_type = NoCustomMask; + + int32_t q_strideM = 0; + int32_t k_strideM = 0; + int32_t v_strideM = 0; + + int32_t o_strideM = 0; + + // Everything below is only used in `advance_to_block` + // and shouldn't use registers + int32_t q_strideH = 0; + int32_t k_strideH = 0; + int32_t v_strideH = 0; + + int64_t q_strideB = 0; + int64_t k_strideB = 0; + int64_t v_strideB = 0; + + int32_t num_batches = 0; + int32_t num_heads = 0; + + // Moves pointers to what we should process + // Returns "false" if there is no work to do + CUTLASS_DEVICE bool advance_to_block() { + auto batch_id = blockIdx.z; + auto head_id = blockIdx.y; + auto query_start = blockIdx.x * kQueriesPerBlock; + + int32_t q_start = 0, k_start = 0; + if (ptr_cumulative_seqlen_Q != nullptr && + ptr_cumulative_seqlen_KV != nullptr) { + // Varlen + ptr_cumulative_seqlen_Q += batch_id; + ptr_cumulative_seqlen_KV += batch_id; + + q_start = ptr_cumulative_seqlen_Q[0]; + k_start = ptr_cumulative_seqlen_KV[0]; + auto q_end = ptr_cumulative_seqlen_Q[1]; + auto k_end = ptr_cumulative_seqlen_KV[1]; + + num_queries = q_end - q_start; + num_keys = k_end - k_start; + + if (query_start >= num_queries) { + return false; + } + + batch_id = 0; + } else { + query_ptr += batch_id * q_strideB; + key_ptr += batch_id * k_strideB; + value_ptr += batch_id * v_strideB; + output_ptr += int64_t(batch_id * num_queries) * o_strideM; + if (output_accum_ptr != nullptr) { + output_accum_ptr += + int64_t(batch_id * num_queries) * (head_dim_value * num_heads); + } + q_start = 0; + k_start = 0; + } + + // Advance to the current batch / head / query_start + query_ptr += (q_start + query_start) * q_strideM + head_id * q_strideH; + key_ptr += k_start * k_strideM + head_id * k_strideH; + + value_ptr += k_start * v_strideM + head_id * v_strideH; + output_ptr += + int64_t(q_start + query_start) * o_strideM + head_id * head_dim_value; + + if (output_accum_ptr != nullptr) { + output_accum_ptr += + int64_t(q_start + query_start) * (head_dim_value * num_heads) + + head_id * head_dim_value; + } else { + // Accumulate directly in the destination buffer (eg for f32) + output_accum_ptr = (accum_t*)output_ptr; + } + + if (logsumexp_ptr != nullptr) { + // lse[batch_id, query_start, head_id] + logsumexp_ptr += batch_id * num_queries * num_heads + head_id + + (q_start + query_start) * num_heads; + } + + // Custom masking + if (custom_mask_type == CausalFromBottomRight) { + causal_diagonal_offset = num_keys - num_queries; + } + if (custom_mask_type == CausalFromTopLeft || + custom_mask_type == CausalFromBottomRight) { + // the bottom row of the current block is query_start + kQueriesPerBlock + // the last active key is then query_start + causal_diagonal_offset + + // kQueriesPerBlock so num_keys is the min between actual num_keys and + // this to avoid extra computations + num_keys = cutlass::fast_min( + int32_t(query_start + causal_diagonal_offset + kQueriesPerBlock), + num_keys); + } + + num_queries -= query_start; + num_batches = 0; // no longer used after + + // If num_queries == 1, and there is only one key head we're wasting + // 15/16th of tensor core compute In that case : + // - we only launch kernels for head_id % kQueriesPerBlock == 0 + // - we iterate over heads instead of queries (strideM = strideH) + if (num_queries == 1 && k_strideH == 0 && v_strideH == 0 && + logsumexp_ptr == nullptr) { + if (head_id % kQueriesPerBlock != 0) { + return false; + } + q_strideM = q_strideH; + num_queries = num_heads; + num_heads = 1; // unused but here for intent + // remove causal since n_query = 1 + // otherwise, offset would change with head ! + custom_mask_type = NoCustomMask; + o_strideM = head_dim_value; + } + + // Make sure the compiler knows these variables are the same on all + // the threads of the warp. + // Only worth doing if they could have been modified above. + query_ptr = warp_uniform(query_ptr); + key_ptr = warp_uniform(key_ptr); + value_ptr = warp_uniform(value_ptr); + output_ptr = warp_uniform(output_ptr); + output_accum_ptr = warp_uniform(output_accum_ptr); + logsumexp_ptr = warp_uniform(logsumexp_ptr); + num_queries = warp_uniform(num_queries); + num_keys = warp_uniform(num_keys); + num_heads = warp_uniform(num_heads); + o_strideM = warp_uniform(o_strideM); + custom_mask_type = warp_uniform(custom_mask_type); + return true; + } + + __host__ dim3 getBlocksGrid() const { + return dim3( + ceil_div(num_queries, (int32_t)kQueriesPerBlock), + num_heads, + num_batches); + } + + __host__ dim3 getThreadsGrid() const { + return dim3(kWarpSize, kNumWarpsPerBlock, 1); + } + }; + + struct MM0 { + /* + In this first matmul, we compute a block of `Q @ K.T`. + While the calculation result is still hot in registers, we update + `mi`, `m_prime`, `s_prime` in shared-memory, and then store this value + into a shared-memory ("AccumulatorSharedStorage") that is used later as + operand A for the second matmul (see MM1) + */ + using GemmType = DefaultGemmType; + + using OpClass = typename GemmType::OpClass; + using DefaultConfig = + typename cutlass::gemm::device::DefaultGemmConfiguration< + OpClass, + ArchTag, + scalar_t, + scalar_t, + scalar_t, // ElementC + accum_t // ElementAccumulator + >; + static constexpr int kAlignmentA = + kIsAligned ? DefaultConfig::kAlignmentA : GemmType::kMinimumAlignment; + static constexpr int kAlignmentB = + kIsAligned ? DefaultConfig::kAlignmentB : GemmType::kMinimumAlignment; + using ThreadblockShape = cutlass::gemm:: + GemmShape; + using WarpShape = cutlass::gemm::GemmShape<32, 32, GemmType::WarpK>; + using DefaultMma = typename cutlass::gemm::threadblock::FindDefaultMma< + scalar_t, // ElementA, + cutlass::layout::RowMajor, // LayoutA, + kAlignmentA, + scalar_t, // ElementB, + cutlass::layout::ColumnMajor, // LayoutB, + kAlignmentB, + accum_t, + cutlass::layout::RowMajor, // LayoutC, + OpClass, + ArchTag, // ArchTag + ThreadblockShape, // ThreadblockShape + WarpShape, // WarpShape + typename GemmType::InstructionShape, // InstructionShape + ArchTag::kMinComputeCapability >= 80 && kIsHalf + ? 4 + : DefaultConfig::kStages, + typename GemmType::Operator // Operator + >::DefaultMma; + using MmaCore = typename DefaultMma::MmaCore; + using IteratorA = typename DefaultMma::IteratorA; + using IteratorB = typename DefaultMma::IteratorB; + using DefaultThreadblockMma = typename DefaultMma::ThreadblockMma; + using Mma = typename cutlass::platform::conditional< + kSingleValueIteration, + typename MakeCustomMma::Mma, + DefaultThreadblockMma>::type; + using AccumLambdaIterator = typename DefaultMmaAccumLambdaIterator< + typename Mma::Operator::IteratorC, + accum_t, + kWarpSize>::Iterator; + static_assert( + MmaCore::WarpCount::kM * MmaCore::WarpCount::kN * + MmaCore::WarpCount::kK == + kNumWarpsPerBlock, + ""); + + // Epilogue to store to shared-memory in a format that we can use later for + // the second matmul + using B2bGemm = typename cutlass::gemm::threadblock::B2bGemm< + typename Mma::Operator::IteratorC, + typename Mma::Operator, + scalar_t, + WarpShape, + ThreadblockShape>; + using AccumulatorSharedStorage = typename B2bGemm::AccumulatorSharedStorage; + }; + + struct MM1 { + /** + Second matmul: perform `attn @ V` where `attn` is the attention (not + normalized) and stored in shared memory + */ + using GemmType = DefaultGemmType; + + using OpClass = typename GemmType::OpClass; + using DefaultConfig = + typename cutlass::gemm::device::DefaultGemmConfiguration< + OpClass, + ArchTag, + scalar_t, + scalar_t, + output_accum_t, // ElementC + accum_t // ElementAccumulator + >; + static constexpr int kAlignmentA = DefaultConfig::kAlignmentA; // from smem + static constexpr int kAlignmentB = + kIsAligned ? DefaultConfig::kAlignmentB : GemmType::kMinimumAlignment; + using ThreadblockShape = cutlass::gemm:: + GemmShape; + using WarpShape = cutlass::gemm::GemmShape<32, 32, GemmType::WarpK>; + using InstructionShape = typename GemmType::InstructionShape; + + using LayoutB = cutlass::layout::RowMajor; + using DefaultGemm = cutlass::gemm::kernel::DefaultGemm< + scalar_t, // ElementA, + cutlass::layout::RowMajor, // LayoutA, + kAlignmentA, + scalar_t, // ElementB, + LayoutB, // LayoutB, + kAlignmentB, + output_accum_t, + cutlass::layout::RowMajor, // LayoutC, + accum_t, + OpClass, + ArchTag, + ThreadblockShape, + WarpShape, + typename GemmType::InstructionShape, + typename DefaultConfig::EpilogueOutputOp, + void, // ThreadblockSwizzle - not used + ArchTag::kMinComputeCapability >= 80 && kIsHalf + ? 4 + : DefaultConfig::kStages, + false, // SplitKSerial + typename GemmType::Operator>; + + using WarpIteratorA = typename cutlass::gemm::threadblock:: + DefaultWarpIteratorAFromSharedMemory< + typename DefaultGemm::Mma::Policy::Operator::Shape, // WarpShape + typename DefaultGemm::Mma::Policy::Operator::InstructionShape, + typename DefaultGemm::Mma::Policy::Operator::IteratorA, + typename DefaultGemm::Mma::Policy>::WarpIterator; + using DefaultMmaFromSmem = + typename cutlass::gemm::threadblock::DefaultMmaFromSharedMemory< + typename DefaultGemm::Mma, + MM0::AccumulatorSharedStorage::Shape::kN, // kMaxK + WarpIteratorA, + false>; // kScaleOperandA + using Mma = typename DefaultMmaFromSmem::Mma; + using IteratorB = typename Mma::IteratorB; + using WarpCount = typename Mma::WarpCount; + static_assert( + WarpCount::kM * WarpCount::kN * WarpCount::kK == kNumWarpsPerBlock, + ""); + + using DefaultEpilogue = typename DefaultGemm::Epilogue; + using OutputTileIterator = + typename cutlass::epilogue::threadblock::PredicatedTileIterator< + typename DefaultEpilogue::OutputTileIterator::ThreadMap, + output_t>; + using OutputTileIteratorAccum = + typename cutlass::epilogue::threadblock::PredicatedTileIterator< + typename DefaultEpilogue::OutputTileIterator::ThreadMap, + output_accum_t>; + }; + + static constexpr int64_t kAlignmentQ = MM0::kAlignmentA; + static constexpr int64_t kAlignmentK = MM0::kAlignmentB; + static constexpr int64_t kAlignmentV = 1; + + // Shared storage - depends on kernel params + struct ScalingCoefs { + cutlass::Array m_prime; + cutlass::Array s_prime; + cutlass::Array mi; + cutlass::Array out_rescale; + cutlass::Array + addition_storage; + }; + + struct SharedStorageEpilogueAtEnd : ScalingCoefs { + struct SharedStorageAfterMM0 { + // Everything here might be overwritten during MM0 + union { + typename MM0::AccumulatorSharedStorage si; + }; + typename MM1::Mma::SharedStorage mm1; + }; + + union { + typename MM0::Mma::SharedStorage mm0; + SharedStorageAfterMM0 after_mm0; + typename MM1::DefaultEpilogue::SharedStorage epilogue; + }; + + CUTLASS_DEVICE typename MM1::DefaultEpilogue::SharedStorage& + epilogue_shared_storage() { + return epilogue; + } + }; + + struct SharedStorageEpilogueInLoop : ScalingCoefs { + struct SharedStorageAfterMM0 { + // Everything here might be overwritten during MM0 + union { + typename MM0::AccumulatorSharedStorage si; + }; + typename MM1::Mma::SharedStorage mm1; + typename MM1::DefaultEpilogue::SharedStorage epilogue; + }; + + union { + typename MM0::Mma::SharedStorage mm0; + SharedStorageAfterMM0 after_mm0; + }; + + CUTLASS_DEVICE typename MM1::DefaultEpilogue::SharedStorage& + epilogue_shared_storage() { + return after_mm0.epilogue; + } + }; + + using SharedStorage = typename cutlass::platform::conditional< + kSingleValueIteration || kKeepOutputInRF, + SharedStorageEpilogueAtEnd, + SharedStorageEpilogueInLoop>::type; + + static bool __host__ check_supported(Params const& p) { + CHECK_ALIGNED_PTR(p.query_ptr, kAlignmentQ); + CHECK_ALIGNED_PTR(p.key_ptr, kAlignmentK); + CHECK_ALIGNED_PTR(p.value_ptr, kAlignmentV); + NATTEN_CHECK( + p.q_strideM % kAlignmentQ == 0, + "query is not correctly aligned (strideM)"); + NATTEN_CHECK( + p.k_strideM % kAlignmentK == 0, + "key is not correctly aligned (strideM)"); + NATTEN_CHECK( + p.v_strideM % kAlignmentV == 0, + "value is not correctly aligned (strideM)"); + NATTEN_CHECK( + p.num_heads <= 1 || p.q_strideH % kAlignmentQ == 0, + "query is not correctly aligned (strideH)"); + NATTEN_CHECK( + p.num_heads <= 1 || p.k_strideH % kAlignmentK == 0, + "key is not correctly aligned (strideH)"); + NATTEN_CHECK( + p.num_heads <= 1 || p.v_strideH % kAlignmentV == 0, + "value is not correctly aligned (strideH)"); + NATTEN_CHECK( + p.custom_mask_type < NumCustomMaskTypes, + "invalid value for `custom_mask_type`"); + return true; + } + + static void CUTLASS_DEVICE attention_kernel(Params& p) { + // In this block, we will only ever: + // - read query[query_start:query_end, :] + // - write to output[query_start:query_end, :] + + extern __shared__ char smem_buffer[]; + SharedStorage& shared_storage = *((SharedStorage*)smem_buffer); + auto& m_prime = shared_storage.m_prime; + auto& s_prime = shared_storage.s_prime; + auto& mi = shared_storage.mi; + auto& out_rescale = shared_storage.out_rescale; + const uint32_t query_start = blockIdx.x * kQueriesPerBlock; + + static_assert(kQueriesPerBlock < kNumWarpsPerBlock * kWarpSize, ""); + if (thread_id() < kQueriesPerBlock) { + s_prime[thread_id()] = accum_t(0); + out_rescale[thread_id()] = accum_t(1.0); + m_prime[thread_id()] = + -cutlass::platform::numeric_limits::infinity(); + mi[thread_id()] = -cutlass::platform::numeric_limits::infinity(); + } + typename MM1::Mma::FragmentC accum_o; + accum_o.clear(); + + auto createOutputIter = [&](int col) -> typename MM1::OutputTileIterator { + using OutputTileIterator = typename MM1::OutputTileIterator; + return OutputTileIterator( + typename OutputTileIterator::Params{(int32_t)p.o_strideM}, + p.output_ptr, + typename OutputTileIterator::TensorCoord{ + p.num_queries, p.head_dim_value}, + thread_id(), + {0, col}); + }; + + auto createOutputAccumIter = [&](int col) -> + typename MM1::OutputTileIteratorAccum { + using OutputTileIteratorAccum = typename MM1::OutputTileIteratorAccum; + return OutputTileIteratorAccum( + typename OutputTileIteratorAccum::Params{ + (int32_t)(p.head_dim_value * p.num_heads)}, + p.output_accum_ptr, + typename OutputTileIteratorAccum::TensorCoord{ + p.num_queries, p.head_dim_value}, + thread_id(), + {0, col}); + }; + + // Iterate through keys + for (int32_t iter_key_start = 0; iter_key_start < p.num_keys; + iter_key_start += kKeysPerBlock) { + int32_t problem_size_0_m = + cutlass::fast_min((int32_t)kQueriesPerBlock, p.num_queries); + int32_t problem_size_0_n = cutlass::fast_min( + int32_t(kKeysPerBlock), p.num_keys - iter_key_start); + int32_t const& problem_size_0_k = p.head_dim; + int32_t const& problem_size_1_n = p.head_dim_value; + int32_t const& problem_size_1_k = problem_size_0_n; + + auto prologueV = [&](int blockN) { + typename MM1::Mma::IteratorB iterator_V( + typename MM1::IteratorB::Params{typename MM1::LayoutB(p.v_strideM)}, + const_cast(p.value_ptr + iter_key_start * p.v_strideM), + {problem_size_1_k, problem_size_1_n}, + thread_id(), + cutlass::MatrixCoord{0, blockN * MM1::Mma::Shape::kN}); + MM1::Mma::prologue( + shared_storage.after_mm0.mm1, + iterator_V, + thread_id(), + problem_size_1_k); + }; + + __syncthreads(); // Need to have shared memory initialized, and `m_prime` + // updated from end of prev iter + // + // MATMUL: Q.K_t + // + // Computes the block-matrix product of: + // (a) query[query_start:query_end, :] + // with + // (b) key[iter_key_start:iter_key_start + kKeysPerBlock] + // and stores that into `shared_storage.si` + // + + // Compute threadblock location + cutlass::gemm::GemmCoord tb_tile_offset = {0, 0, 0}; + + cutlass::MatrixCoord tb_offset_A{ + tb_tile_offset.m() * MM0::Mma::Shape::kM, tb_tile_offset.k()}; + + cutlass::MatrixCoord tb_offset_B{ + tb_tile_offset.k(), tb_tile_offset.n() * MM0::Mma::Shape::kN}; + + // Construct iterators to A and B operands + typename MM0::IteratorA iterator_A( + typename MM0::IteratorA::Params( + typename MM0::MmaCore::LayoutA(p.q_strideM)), + const_cast(p.query_ptr), + {problem_size_0_m, problem_size_0_k}, + thread_id(), + tb_offset_A); + + typename MM0::IteratorB iterator_B( + typename MM0::IteratorB::Params( + typename MM0::MmaCore::LayoutB(p.k_strideM)), + const_cast(p.key_ptr + iter_key_start * p.k_strideM), + {problem_size_0_k, problem_size_0_n}, + thread_id(), + tb_offset_B); + + auto my_warp_id = warp_uniform(warp_id()); + auto my_lane_id = lane_id(); + + // Construct thread-scoped matrix multiply + typename MM0::Mma mma( + shared_storage.mm0, thread_id(), my_warp_id, my_lane_id); + + typename MM0::Mma::FragmentC accum; + + accum.clear(); + + auto gemm_k_iterations = + (problem_size_0_k + MM0::Mma::Shape::kK - 1) / MM0::Mma::Shape::kK; + + // Compute threadblock-scoped matrix multiply-add + mma(gemm_k_iterations, accum, iterator_A, iterator_B, accum); + __syncthreads(); + + if (kPreloadV) { + prologueV(0); + } else { + MM1::Mma::drain_cp_asyncs(); + } + + typename MM0::Mma::Operator::IteratorC::TensorCoord + iteratorC_tile_offset = { + (tb_tile_offset.m() * MM0::Mma::WarpCount::kM) + + (my_warp_id % MM0::Mma::WarpCount::kM), + (tb_tile_offset.n() * MM0::Mma::WarpCount::kN) + + (my_warp_id / MM0::Mma::WarpCount::kM)}; + + // Mask out last if causal + // This is only needed if upper-right corner of current query / key block + // intersects the mask Coordinates of upper-right corner of current block + // is y=query_start x=min(iter_key_start + kKeysPerBlock, num_keys)) The + // first masked element is x = y + offset -> query_start + offset There is + // intersection (and we need to mask) if min(iter_key_start + + // kKeysPerBlock, num_keys)) >= query_start + offset + if (p.custom_mask_type && + cutlass::fast_min(iter_key_start + kKeysPerBlock, p.num_keys) >= + (query_start + p.causal_diagonal_offset)) { + auto query_start = blockIdx.x * kQueriesPerBlock; + auto lane_offset = MM0::AccumLambdaIterator::get_lane_offset( + my_lane_id, my_warp_id, iteratorC_tile_offset); + int32_t last_col; + MM0::AccumLambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) { + // last absolute col is (last absolute query + offset) + // last local col is (last absolute query + offset - + // iter_key_start) + last_col = query_start + accum_m + p.causal_diagonal_offset - + iter_key_start; + }, + [&](int accum_m, int accum_n, int idx) { + if (accum_n > last_col) { + accum[idx] = + -cutlass::platform::numeric_limits::infinity(); + } + }, + [&](int accum_m) {}); + } + + // Update `mi` from accum stored in registers + // Also does accum[i] <- exp(accum[i] - mi) + iterative_softmax( + accum_o, + accum, + mi, + m_prime, + s_prime, + out_rescale, + shared_storage.addition_storage, + my_lane_id, + thread_id(), + my_warp_id, + p.num_keys - iter_key_start, + iter_key_start == 0, + iteratorC_tile_offset, + p.scale); + + // Output results to shared-memory + int warp_idx_mn_0 = my_warp_id % + (MM0::Mma::Base::WarpCount::kM * MM0::Mma::Base::WarpCount::kN); + auto output_tile_coords = cutlass::MatrixCoord{ + warp_idx_mn_0 % MM0::Mma::Base::WarpCount::kM, + warp_idx_mn_0 / MM0::Mma::Base::WarpCount::kM}; + + MM0::B2bGemm::accumToSmem( + shared_storage.after_mm0.si, accum, my_lane_id, output_tile_coords); + + __syncthreads(); + + // + // MATMUL: Attn . V + // Run the matmul `attn @ V` for a block of attn and V. + // `attn` is read from shared memory (in `shared_storage_si`) + // `V` is read from global memory (with iterator_B) + // + + const int64_t nBlockN = kSingleValueIteration + ? 1 + : ceil_div( + (int64_t)problem_size_1_n, int64_t(MM1::ThreadblockShape::kN)); + for (int blockN = 0; blockN < nBlockN; ++blockN) { + int gemm_k_iterations = + (problem_size_1_k + MM1::Mma::Shape::kK - 1) / MM1::Mma::Shape::kK; + + // Compute threadblock-scoped matrix multiply-add and store it in accum + // (in registers) + if (!kPreloadV) { + __syncthreads(); // we share shmem between mma and epilogue + } + + typename MM1::Mma::IteratorB iterator_V( + typename MM1::IteratorB::Params{typename MM1::LayoutB(p.v_strideM)}, + const_cast(p.value_ptr + iter_key_start * p.v_strideM), + {problem_size_1_k, problem_size_1_n}, + thread_id(), + cutlass::MatrixCoord{0, blockN * MM1::Mma::Shape::kN}); + typename MM1::Mma mma_pv( + // operand A: Pij_dropped in shared memory + shared_storage.after_mm0.si.accum_ref(), + // operand B: shared memory staging area for Vj, which is loaded + // from global memory + shared_storage.after_mm0.mm1.operand_B_ref(), + (int)thread_id(), + (int)my_warp_id, + (int)my_lane_id); + mma_pv.set_prologue_done(kPreloadV); + if (!kKeepOutputInRF) { + accum_o.clear(); + } + mma_pv(gemm_k_iterations, accum_o, iterator_V, accum_o); + __syncthreads(); + + if (kPreloadV && !kSingleValueIteration && blockN + 1 < nBlockN) { + prologueV(blockN + 1); + } + + if (!kKeepOutputInRF) { + int first_key = 0; + + // int first_key_block = 0; + MM1::Mma::drain_cp_asyncs(); + DISPATCH_BOOL( + iter_key_start == first_key, kIsFirst, ([&] { + DISPATCH_BOOL( + (iter_key_start + kKeysPerBlock) >= p.num_keys, + kIsLast, + ([&] { + using DefaultEpilogue = typename MM1::DefaultEpilogue; + using DefaultOp = + typename MM1::DefaultConfig::EpilogueOutputOp; + using ElementCompute = typename DefaultOp::ElementCompute; + using EpilogueOutputOp = typename cutlass::epilogue:: + thread::MemoryEfficientAttentionNormalize< + typename cutlass::platform::conditional< + kIsLast, + output_t, + output_accum_t>::type, + output_accum_t, + DefaultOp::kCount, + typename DefaultOp::ElementAccumulator, + ElementCompute, + kIsFirst, + kIsLast, + cutlass::Array>; + using Epilogue = typename cutlass::epilogue::threadblock:: + EpiloguePipelined< + typename DefaultEpilogue::Shape, + typename MM1::Mma::Operator, + DefaultEpilogue::kPartitionsK, + typename cutlass::platform::conditional< + kIsLast, + typename MM1::OutputTileIterator, + typename MM1::OutputTileIteratorAccum>::type, + typename DefaultEpilogue:: + AccumulatorFragmentIterator, + typename DefaultEpilogue::WarpTileIterator, + typename DefaultEpilogue::SharedLoadIterator, + EpilogueOutputOp, + typename DefaultEpilogue::Padding, + DefaultEpilogue::kFragmentsPerIteration, + true, // IterationsUnroll + typename MM1::OutputTileIteratorAccum // Read + // iterator + >; + + int col = blockN * MM1::Mma::Shape::kN; + auto source_iter = createOutputAccumIter(col); + auto dest_iter = call_conditional< + kIsLast, + decltype(createOutputIter), + decltype(createOutputAccumIter)>:: + apply(createOutputIter, createOutputAccumIter, col); + EpilogueOutputOp rescale(s_prime, out_rescale); + Epilogue epilogue( + shared_storage.epilogue_shared_storage(), + thread_id(), + my_warp_id, + my_lane_id); + epilogue(rescale, dest_iter, accum_o, source_iter); + })); + })); + if (!kSingleValueIteration) { + __syncthreads(); + } + } + } + __syncthreads(); // we modify `m_prime` after + } + + if (kKeepOutputInRF) { + constexpr bool kIsFirst = true; + constexpr bool kIsLast = true; + using DefaultEpilogue = typename MM1::DefaultEpilogue; + using DefaultOp = typename MM1::DefaultConfig::EpilogueOutputOp; + using ElementCompute = typename DefaultOp::ElementCompute; + using EpilogueOutputOp = + typename cutlass::epilogue::thread::MemoryEfficientAttentionNormalize< + output_t, // output + output_accum_t, // source + DefaultOp::kCount, + typename DefaultOp::ElementAccumulator, // accum + output_accum_t, // compute + kIsFirst, + kIsLast, + cutlass::Array>; + using Epilogue = + typename cutlass::epilogue::threadblock::EpiloguePipelined< + typename DefaultEpilogue::Shape, + typename MM1::Mma::Operator, + DefaultEpilogue::kPartitionsK, + typename MM1::OutputTileIterator, // destination + typename DefaultEpilogue::AccumulatorFragmentIterator, + typename DefaultEpilogue::WarpTileIterator, + typename DefaultEpilogue::SharedLoadIterator, + EpilogueOutputOp, + typename DefaultEpilogue::Padding, + DefaultEpilogue::kFragmentsPerIteration, + true, // IterationsUnroll + typename MM1::OutputTileIteratorAccum // source tile + >; + auto dest_iter = createOutputIter(0); + EpilogueOutputOp rescale(s_prime, out_rescale); + Epilogue epilogue( + shared_storage.epilogue_shared_storage(), + thread_id(), + warp_id(), + lane_id()); + MM1::Mma::drain_cp_asyncs(); + epilogue(rescale, dest_iter, accum_o); + } + + // 7. Calculate logsumexp + // To make the backward easier, we pad logsumexp with `inf` + // this avoids a few bound checks, and is not more expensive during fwd + static_assert(kQueriesPerBlock < kNumWarpsPerBlock * kWarpSize, ""); + if (p.logsumexp_ptr && thread_id() < kQueriesPerBlock) { + constexpr float kLog2e = 1.4426950408889634074; // log_2(e) = M_LOG2E + if (thread_id() < p.num_queries) { + auto query_offset = static_cast(thread_id()) * p.num_heads; + p.logsumexp_ptr[query_offset] = accum_t(mi[thread_id()] / kLog2e) + + cutlass::fast_log(accum_t(s_prime[thread_id()])); + } + } + } + + template + CUTLASS_DEVICE static void iterative_softmax( + typename WarpIteratorC::Fragment& frag_o, // output so far + typename WarpIteratorC::Fragment& frag, + cutlass::Array& mi, + cutlass::Array& m_prime, + cutlass::Array& s_prime, + cutlass::Array& out_rescale, + cutlass::Array& + addition_storage, + int8_t lane_id, + int8_t thread_id, + int8_t warp_id, + int max_col, + bool is_first, + typename WarpIteratorC::TensorCoord const& tile_offset, + float scaling) { + /* Iterates on the accumulator and corresponding position on result matrix + + (1) Update `mi[r]` to the max value of the row `r` + (2) In a second iteration do the following: + (a) accum <- exp(accum - mi) + (b) m_prime <- exp(m_prime - mi) + (c) s_prime <- s_prime * m_prime + sum(accum) + + All of this is done on registers, before we store all of this + on shared memory for the next matmul with Value. + */ + using Fragment = typename WarpIteratorC::Fragment; + using LambdaIterator = typename DefaultMmaAccumLambdaIterator< + WarpIteratorC, + accum_t, + kWarpSize>::Iterator; + // Convert to `accum_t` (rather than double) + constexpr float kLog2e = 1.4426950408889634074; // log_2(e) = M_LOG2E + + static_assert(kQueriesPerBlock % kNumWarpsPerBlock == 0, ""); + static constexpr int kLinesPerWarp = kQueriesPerBlock / kNumWarpsPerBlock; + + frag = cutlass::multiplies()(scaling * kLog2e, frag); + + auto lane_offset = + LambdaIterator::get_lane_offset(lane_id, warp_id, tile_offset); + + // First update `mi` to the max per-row + { + accum_t max; + LambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) { + max = -cutlass::platform::numeric_limits::infinity(); + }, + [&](int accum_m, int accum_n, int idx) { + if (accum_n < max_col) { + max = cutlass::fast_max(max, frag[idx]); + } + }, + [&](int accum_m) { + // Having 4x atomicMax seems faster than reduce within warp + // first... + atomicMaxFloat(&mi[accum_m], max); + }); + } + + // Make sure we all share the update values for `mi` + __syncthreads(); + + // Doing this `exp` is quite expensive. Let's + // split it across the warps + bool restore_mi_to_minus_inf = false; + if (lane_id < kLinesPerWarp) { + int id = warp_id * kLinesPerWarp + lane_id; + auto m_prime_id = m_prime[id]; + auto mi_id = mi[id]; + bool changed = m_prime_id < mi_id; // `false` if both are -inf + if (changed) { + auto m_prime_exp = exp2f(m_prime_id - mi_id); + out_rescale[id] = m_prime_exp; + s_prime[id] *= m_prime_exp; + } else { + out_rescale[id] = 1.0f; + } + } + __syncthreads(); // Update output fragments + if (kKeepOutputInRF && !is_first) { + accum_t line_rescale; + LambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) { line_rescale = out_rescale[accum_m]; }, + [&](int accum_m, int accum_n, int idx) { + frag_o[idx] = frag_o[idx] * line_rescale; + }, + [&](int accum_m) {}); + } + // Update accum_m, accum_n, ... + { + accum_t mi_row, total_row; + LambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) { mi_row = mi[accum_m]; }, + [&](int accum_m, int accum_n, int idx) { + frag[idx] = + (accum_n < max_col) ? exp2f(frag[idx] - mi_row) : accum_t(0.0); + }, + [&](int accum_m) {}); + LambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) { total_row = 0.0; }, + [&](int accum_m, int accum_n, int idx) { total_row += frag[idx]; }, + [&](int accum_m) { + if (LambdaIterator::reduceSameRow( + lane_id, total_row, [](accum_t a, accum_t b) { + return a + b; + })) { + // NOTE: we could atomically add `total_row` to `s_prime`, but + // it's faster (and deterministic) to avoid atomics here + addition_storage + [accum_m + kQueriesPerBlock * tile_offset.column()] = + total_row; + } + }); + } + __syncthreads(); + if (lane_id < kLinesPerWarp) { + int id = warp_id * kLinesPerWarp + lane_id; + accum_t total_row = s_prime[id]; + if (restore_mi_to_minus_inf) { + // Restore `mi`, see above when we set `restore_mi_to_minus_inf=true` + mi[id] = -cutlass::platform::numeric_limits::infinity(); + } else { + m_prime[id] = mi[id]; + } + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < MM0::MmaCore::WarpCount::kN; ++i) { + total_row += addition_storage[id + kQueriesPerBlock * i]; + } + s_prime[id] = total_row; + } + } + + static CUTLASS_DEVICE int8_t lane_id() { + return threadIdx.x; + } + static CUTLASS_DEVICE int8_t warp_id() { + return threadIdx.y; + } + static CUTLASS_DEVICE int16_t thread_id() { + return threadIdx.x + threadIdx.y * blockDim.x; + } +}; + +} // namespace fmha +} // namespace cuda +} // namespace natten diff --git a/natten/csrc/include/natten/cuda/fmha/transform/tile_smem_loader.h b/natten/csrc/include/natten/cuda/fmha/transform/tile_smem_loader.h new file mode 100644 index 00000000..10bbfe3e --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha/transform/tile_smem_loader.h @@ -0,0 +1,66 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +template < + typename scalar_t, // scalar type + typename ThreadblockTileShape, // size of tile to load + int Threads, // number of participating threads + int ElementsPerAccess> // thread access width in elements +class TileSmemLoader { + public: + using SmemTile = + cutlass::AlignedBuffer; + + using ThreadMap = cutlass::transform::PitchLinearStripminedThreadMap< + cutlass::layout::PitchLinearShape< + ThreadblockTileShape::kColumn, // contiguous + ThreadblockTileShape::kRow>, // strided + Threads, // Threads + ElementsPerAccess>; // ElementsPerAccess + + using GmemTileIterator = + cutlass::transform::threadblock::PredicatedTileIterator< + ThreadblockTileShape, // Shape + scalar_t, // Element + cutlass::layout::RowMajor, // Layout + 0, // AdvanceRank + ThreadMap>; // ThreadMap + + using SmemTileIterator = cutlass::transform::threadblock::RegularTileIterator< + ThreadblockTileShape, // Shape + scalar_t, // Element + cutlass::layout::RowMajor, // Layout + 0, // AdvanceRank + ThreadMap>; // ThreadMap + + using Fragment = typename GmemTileIterator::Fragment; + + /// load a tile from global memory into shared memory + CUTLASS_DEVICE + static void load( + GmemTileIterator tile_load_iter, + SmemTileIterator tile_store_iter) { + Fragment tb_frag; + tb_frag.clear(); + tile_load_iter.load(tb_frag); + tile_store_iter.store(tb_frag); + + __syncthreads(); + } +}; diff --git a/natten/csrc/include/natten/cuda/fmha_blackwell/collective/fmha_common.hpp b/natten/csrc/include/natten/cuda/fmha_blackwell/collective/fmha_common.hpp new file mode 100644 index 00000000..6c96030d --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_blackwell/collective/fmha_common.hpp @@ -0,0 +1,180 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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 +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 < + class a_type, + class b_type, + class c_type, + int M, + int N, + UMMA::Major a_major, + UMMA::Major b_major, + UMMA::ScaleIn a_neg, + UMMA::ScaleIn b_neg, + class... TAs, + class... TMs> +CUTE_HOST_DEVICE constexpr auto to_tiled_mma_sm100_ts( + TiledMMA< + MMA_Atom< + MMA_Traits< + SM100_MMA_F8F6F4_SS, + a_type, + b_type, + c_type, + cute::C, + cute::C, + cute::integral_constant, + cute::integral_constant, + cute::integral_constant, + cute::integral_constant>, + TAs...>, + TMs...>) { + return TiledMMA< + MMA_Atom< + MMA_Traits>, + TAs...>, + TMs...>{}; +} + +template < + class a_type, + class b_type, + class c_type, + int M, + int N, + UMMA::Major a_major, + UMMA::Major b_major, + UMMA::ScaleIn a_neg, + UMMA::ScaleIn b_neg, + class... TAs, + class... TMs> +CUTE_HOST_DEVICE constexpr auto to_tiled_mma_sm100_ts(TiledMMA< + MMA_Atom< + SM100_MMA_F16BF16_SS< + a_type, + b_type, + c_type, + M, + N, + a_major, + b_major, + a_neg, + b_neg>, + TAs...>, + TMs...>) { + return TiledMMA< + MMA_Atom< + SM100_MMA_F16BF16_TS< + a_type, + b_type, + c_type, + M, + N, + a_major, + b_major, + a_neg, + b_neg, + UMMA::Saturate::False>, + 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/natten/csrc/include/natten/cuda/fmha_blackwell/collective/fmha_fusion.hpp b/natten/csrc/include/natten/cuda/fmha_blackwell/collective/fmha_fusion.hpp new file mode 100644 index 00000000..d2a9a72b --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_blackwell/collective/fmha_fusion.hpp @@ -0,0 +1,379 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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 ResidualMaskForBackward : 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 (!elem_less(pos, select<0, 1>(problem_size))) { + acc_qk(i) = -INFINITY; + } + } + } +}; + +// There are two ways to do causal if N_Q != N_K +// (1) The Q is at the beginning of the matrix +// (2) The Q is at the end of the matrix +template +struct CausalMask : NoMask { + using Base = NoMask; + + static constexpr bool IsQBegin = kIsQBegin; + + 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 max_blocks_k = + Base::get_trip_count(blk_coord, tile_shape, problem_size); + if constexpr (IsQBegin) { + int max_blocks_q = ceil_div( + (get<0>(blk_coord) + 1) * get<0>(tile_shape), get<1>(tile_shape)); + return std::min(max_blocks_k, max_blocks_q); + } else { + const int offset_q = get<1>(problem_size) - get<0>(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) { + int trip_count = get_trip_count(blk_coord, tile_shape, problem_size); + if constexpr (IsQBegin) { + return std::min( + trip_count, int(ceil_div(size<0>(tile_shape), size<1>(tile_shape)))); + } else { + const int offset_tile_q = + (get<1>(problem_size) - get<0>(problem_size)) % get<1>(tile_shape); + return std::min( + trip_count, + int(ceil_div( + get<0>(tile_shape) + offset_tile_q, get<1>(tile_shape)))); + } + } + + 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 the default setting. + // (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 set kIsQBegin=false + + if constexpr (IsQBegin) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(acc_qk); i++) { + auto pos = index_qk(i); + if ((get<0>(pos) < get<1>(pos)) || + (get<1>(pos) >= get<1>(problem_size))) { + acc_qk(i) = -INFINITY; + } + } + } else { + const auto offset_q = get<1>(problem_size) - 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; + } + } + } + } +}; + +template +struct CausalForBackwardMask : CausalMask, ResidualMaskForBackward { + using Base = CausalMask; + + 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 = 0; + if constexpr (!kIsQBegin) { + offset_q = get<1>(problem_size) - get<0>(problem_size); + } + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(acc_qk); i++) { + auto pos = index_qk(i); + bool masked = (get<0>(pos) + offset_q < get<1>(pos)) || + !elem_less(pos, problem_size); + if (masked) { + acc_qk(i) = -INFINITY; + } + } + } +}; + +struct VariableLength { + int max_length; + int* cumulative_length = nullptr; + int total_length = -1; + + CUTE_HOST_DEVICE operator int() const { + return max_length; + } +}; + +template +struct is_variable_length_impl : std::false_type {}; +template <> +struct is_variable_length_impl : std::true_type {}; +template +constexpr bool is_variable_length_v = + is_variable_length_impl>::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.cumulative_length[idx + 1] - s.cumulative_length[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.cumulative_length[idx]); + } else { + return c; + } + }); + return cute::make_tuple(new_shape, new_coord); +} + +template +CUTE_HOST_DEVICE constexpr auto apply_variable_length_offset( + Shape const& shape, + Coord const& coord) { + auto idx = back(back(coord)); + auto result_shape = transform_leaf(shape, [&](auto const& s) { + if constexpr (is_variable_length_v) { + return s.cumulative_length[idx + 1] - s.cumulative_length[idx]; + } else { + return s; + } + }); + auto result_offset = + transform_leaf(coord, shape, [&](auto const& c, auto const& s) { + if constexpr (is_variable_length_v) { + return s.cumulative_length[idx]; + } else { + return _0{}; + } + }); + return cute::make_tuple(result_shape, result_offset); +} + +} // namespace cutlass::fmha::collective + +namespace cute { + +template <> +struct is_integral : true_type {}; + +CUTE_HOST_DEVICE +void print(cutlass::fmha::collective::VariableLength a) { + printf("Varlen<%d, %p>", a.max_length, a.cumulative_length); +} + +} // namespace cute diff --git a/natten/csrc/include/natten/cuda/fmha_blackwell/collective/sm100_fmha_fwd_epilogue_tma_warpspecialized.hpp b/natten/csrc/include/natten/cuda/fmha_blackwell/collective/sm100_fmha_fwd_epilogue_tma_warpspecialized.hpp new file mode 100644 index 00000000..73f7825d --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_blackwell/collective/sm100_fmha_fwd_epilogue_tma_warpspecialized.hpp @@ -0,0 +1,240 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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" + +namespace cutlass::fmha::collective { + +template < + class Element, + class ElementAcc, + class TileShape, // Q, D, _ + class StrideO, // Q, D, B + class StrideLSE_, // Q, B + class OrderLoadEpilogue = cute::false_type> +struct Sm100FmhaFwdEpilogueTmaWarpspecialized { + using Pipeline = cutlass::PipelineAsync<2>; + + // 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; + using StrideLSE = StrideLSE_; + using ElementOut = Element; + + static const int NumWarpsEpilogue = 1; + static const int NumWarpsLoad = 1; + + struct TensorStorage { + using SmemLayoutO = SmemLayoutO_; + cute::array_aligned> smem_o; + }; + + struct Arguments { + Element* ptr_O; + StrideO dO; + + ElementAcc* ptr_LSE; + StrideLSE dLSE; + }; + + 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; + + ElementAcc* ptr_LSE; + StrideLSE dLSE; + }; + + // FMHA and MLA have different input ProblemShapes; + // get problem_shape_O according to the input ProblemShape. + template + CUTLASS_DEVICE static constexpr auto get_problem_shape_O( + ProblemShape const& problem_shape) { + if constexpr (rank_v(ProblemShape{}))> == 2) { + return replace<1>( + select<0, 2, 3>(problem_shape), get<2, 0>(problem_shape)); + } else { + return select<0, 2, 3>(problem_shape); + } + } + + template + static Params to_underlying_arguments( + ProblemShape const& problem_shape, + Arguments const& args, + void* workspace = nullptr) { + auto ptr_O = args.ptr_O; + StrideO dO = args.dO; + + auto problem_shape_O = get_problem_shape_O(problem_shape); + + if constexpr (is_variable_length_v>) { + auto cumulative_length_q = get<0>(problem_shape).cumulative_length; + if (cumulative_length_q != nullptr) { + int max_length_q = get<0>(problem_shape).max_length; + // for variable sequence lenght, the batch is in units of row_stride + get<2, 1>(dO) = get<0>(dO); + get<2, 1>(problem_shape_O) = + max_length_q * (1 + get<2, 1>(problem_shape_O)); + // offset ptr by the amount we add back in later + ptr_O -= max_length_q * get<0>(dO); + } + } + + auto tma_store_o = make_tma_copy( + SM90_TMA_STORE{}, + make_tensor(ptr_O, problem_shape_O, dO), + SmemLayoutO{}(_, _, _0{})); + + return {tma_store_o, args.ptr_LSE, args.dLSE}; + } + + 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_in, + ProblemShape const& problem_shape, + Params const& params, + ParamsProblemShape const& params_problem_shape, + TensorStorage& shared_storage, + Pipeline& pipeline, + typename Pipeline::PipelineState& pipeline_consumer_state) { + BlkCoord blk_coord = blk_coord_in; + 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; + + Tensor mO_qdl_p = + params.tma_store_o.get_tma_tensor(get_problem_shape_O(problem_shape)); + // offset mode 0 by (max_length - real_length) + // offset mode 3,1 by cumulative_length + real_length + // the ptr is already offset by - max_length + // so in total this achieves + int offs_0 = 0; + int offs_2_1 = 0; + + if constexpr (is_variable_length_v< + tuple_element_t<0, ParamsProblemShape>>) { + auto cumulative_length_q = get<0>(params_problem_shape).cumulative_length; + if (cumulative_length_q != nullptr) { + int max_length_q = get<0>(params_problem_shape).max_length; + offs_0 = max_length_q - get<0>(problem_shape); + offs_2_1 = + cumulative_length_q[get<2, 1>(blk_coord)] + get<0>(problem_shape); + get<2, 1>(blk_coord) = 0; + } + } + + Tensor mO_qdl = domain_offset( + make_coord(offs_0, _0{}, make_coord(_0{}, offs_2_1)), mO_qdl_p); + + Tensor gO_qdl = + local_tile(mO_qdl, TileShape{}, make_coord(_, _, _), Step<_1, _1, X>{}); + Tensor gO = gO_qdl(_, _, _, _0{}, get<2>(blk_coord)); + 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>(); + + if constexpr (cute::is_same_v) { + cutlass::arch::NamedBarrier::arrive( + (NumWarpsLoad + NumWarpsEpilogue) * NumThreadsPerWarp, + cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); + } + + pipeline.consumer_release(pipeline_release_state); + ++pipeline_release_state; + } +}; + +} // namespace cutlass::fmha::collective diff --git a/natten/csrc/include/natten/cuda/fmha_blackwell/collective/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp b/natten/csrc/include/natten/cuda/fmha_blackwell/collective/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp new file mode 100644 index 00000000..e5014e3e --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_blackwell/collective/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp @@ -0,0 +1,1387 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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 "natten/cuda/fmha_blackwell/collective/fmha_common.hpp" +#include "natten/cuda/fmha_blackwell/collective/fmha_fusion.hpp" +#include "natten/cuda/fmha_blackwell/collective/sm100_fmha_load_tma_warpspecialized.hpp" + +namespace cutlass::fmha::collective { + +using namespace cute; + +template < + class Element_, + class ElementQK_, + class ElementPV_, + class TileShape_, + class StrideQ_, + class StrideK_, + class StrideV_, + class Mask_, + // shape here is QG K H + // and referes to the two softmax warps + // (2, 1, 1) means that they are stacked (best for large Q since it loads + // the least K/V) (1, 2, 1) means they sit side by side (best for small Q / + // large K) + class ThreadShape = Shape<_2, _1, _1>, + // Since shared memory is sufficient for FMHA, there is no need to reuse + // shared memory. + class OrderLoadEpilogue = cute::false_type> +struct Sm100FmhaFwdMainloopTmaWarpspecialized { + using Element = Element_; + using ElementQK = ElementQK_; + using ElementPV = ElementPV_; + using TileShape = TileShape_; + using StrideQ = StrideQ_; + using StrideK = StrideK_; + using StrideV = StrideV_; + using Mask = Mask_; + + static constexpr int StageCountQ = 2; + static constexpr int StageCountKV = sizeof(Element_) == 1 ? 4 : 3; + + 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{})); + + // Reuse shared memory for V and O. + static constexpr bool IsOrderLoadEpilogue = + std::is_same_v; + 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::PipelineTmaUmmaAsync< + StageCountQ, + typename CollectiveMmaQK::AtomThrShapeMNK>; + + // from load to mma warp, protects k/v in smem + using PipelineKV = cutlass::PipelineTmaUmmaAsync< + StageCountKV, + typename CollectiveMmaQK::AtomThrShapeMNK>; + + // 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( + TransactionBytesLoadK == TransactionBytesLoadV, + "K and V smem layouts must be of equal size"); + + using Load = Sm100FmhaLoadTmaWarpspecialized< + Element, + StrideQ, + StrideK, + StrideV, + CollectiveMmaQK, + CollectiveMmaPV, + SmemLayoutQ, + SmemLayoutK, + SmemLayoutV, + TensorStorage, + PipelineQ, + PipelineKV, + Mask, + TileShape>; + + 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 < + bool need_apply_mask, + class Stage, + class BlkCoord, + class CoordTensor, + class ProblemShape> + CUTLASS_DEVICE auto softmax_step( + float& row_max, + float& row_sum, + Stage stage, + bool final_call, + BlkCoord const& blk_coord, + CoordTensor 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 = + size<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 = (old_row_max == row_max_safe) + ? 0.5f + : 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< + kCorrectionTileSize == 32, + SM100_TMEM_LOAD_32dp32b32x, + SM100_TMEM_LOAD_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 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<2>(TileShape{}) / 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<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))); + + if (scale != 1.0f) { + 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 < + class BlkCoord, + class ProblemShape, + class ParamsProblemShape, + class TensorStorageEpi, + class CollectiveEpilogue> + CUTLASS_DEVICE auto correction( + BlkCoord const& blk_coord, + Params const& params, + ProblemShape const& problem_shape, + ParamsProblemShape const& params_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 = + (tTMEM_LOADVrS(kIdxOldRowMax) == tTMEM_LOADVrS(kIdxNewRowMax)) + ? 1.0f + : ::exp2f( + params.scale_softmax_log2 * + (tTMEM_LOADVrS(kIdxOldRowMax) - tTMEM_LOADVrS(kIdxNewRowMax))); + + pipeline_o.consumer_wait(pipeline_o_consumer_state); + + bool warp_do_correction = __any_sync(0xFFFFFFFF, scale != 1.0f); + if (warp_do_correction) { + 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 = (tTMEM_LOADVrS(kIdxOldRowMax) == tTMEM_LOADVrS(kIdxNewRowMax)) + ? 1.0f + : ::exp2f( + params.scale_softmax_log2 * + (tTMEM_LOADVrS(kIdxOldRowMax) - tTMEM_LOADVrS(kIdxNewRowMax))); + + pipeline_o.consumer_wait(pipeline_o_consumer_state); + + warp_do_correction = __any_sync(0xFFFFFFFF, scale != 1.0f); + if (warp_do_correction) { + 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), + select<0, 3>(problem_shape), + epilogue.params.dLSE); + + correction_epilogue( + params.scale_output / tTMEM_LOADVrS(kIdxFinalRowSum), _0{}, sO); + + if (epilogue.params.ptr_LSE != nullptr) { + int row_idx = + get<0>(tTMEM_LOADVcS(_0{})) + get<0>(TileShape{}) * get<0>(blk_coord); + + int row_offset = 0; + if constexpr (is_variable_length_v< + tuple_element_t<0, ParamsProblemShape>>) { + row_offset = get<0>(params_problem_shape) + .cumulative_length[get<2, 1>(blk_coord)]; + } + + ElementPV lse = cutlass::fast_log(tTMEM_LOADVrS(kIdxFinalRowSum)) + + params.scale_softmax * tTMEM_LOADVrS(kIdxFinalRowMax); + + if (row_idx < get<0>(problem_shape)) { + gLSE(row_idx + row_offset, get<2>(blk_coord)) = lse; + } + } + + 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 row_idx = get<0>(tTMEM_LOADVcS(_0{})) + + get<0>(TileShape{}) * get<0>(blk_coord) + get<0>(TileShapeQK{}); + + ElementPV lse = cutlass::fast_log(tTMEM_LOADVrS(kIdxFinalRowSum)) + + params.scale_softmax * tTMEM_LOADVrS(kIdxFinalRowMax); + + int row_offset = 0; + if constexpr (is_variable_length_v< + tuple_element_t<0, ParamsProblemShape>>) { + row_offset = get<0>(params_problem_shape) + .cumulative_length[get<2, 1>(blk_coord)]; + } + + if (row_idx < get<0>(problem_shape)) { + gLSE(row_idx + row_offset, get<2>(blk_coord)) = lse; + } + } + + 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 < + class BlkCoord, + class ProblemShape, + class ParamsProblemShape, + class TensorStorageEpi, + class CollectiveEpilogue> + 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), + select<0, 3>(problem_shape), + epilogue.params.dLSE); + float lse = -INFINITY; + int thread_idx = threadIdx.x % (4 * NumThreadsPerWarp); + +#define DSHOW(x) \ + print(#x ": "); \ + print(x); \ + print("\n") + if (threadIdx.x % 128 == 0 && block0()) { + DSHOW(sO); + } +#if 1 + + 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{})); +#endif + + if (epilogue.params.ptr_LSE != nullptr) { + int row_idx = thread_idx + get<0>(TileShape{}) * get<0>(blk_coord); + + int row_offset = 0; + if constexpr (is_variable_length_v< + tuple_element_t<0, ParamsProblemShape>>) { + row_offset = get<0>(params_problem_shape) + .cumulative_length[get<2, 1>(blk_coord)]; + } + + if (row_idx < get<0>(problem_shape)) { + gLSE(row_idx + row_offset, get<2>(blk_coord)) = lse; + } + } + + 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 row_idx = thread_idx + get<0>(TileShape{}) * get<0>(blk_coord) + + get<0>(TileShapeQK{}); + + int row_offset = 0; + if constexpr (is_variable_length_v< + tuple_element_t<0, ParamsProblemShape>>) { + row_offset = get<0>(params_problem_shape) + .cumulative_length[get<2, 1>(blk_coord)]; + } + + if (row_idx < get<0>(problem_shape)) { + gLSE(row_idx + row_offset, get<2>(blk_coord)) = lse; + } + } + + 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/natten/csrc/include/natten/cuda/fmha_blackwell/collective/sm100_fmha_load_tma_warpspecialized.hpp b/natten/csrc/include/natten/cuda/fmha_blackwell/collective/sm100_fmha_load_tma_warpspecialized.hpp new file mode 100644 index 00000000..3e37ed24 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_blackwell/collective/sm100_fmha_load_tma_warpspecialized.hpp @@ -0,0 +1,346 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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 "cutlass/arch/memory_sm80.h" +#include "cutlass/cutlass.h" +#include "cutlass/gemm/collective/collective_builder.hpp" + +#include "natten/cuda/fmha_blackwell/collective/fmha_common.hpp" +#include "natten/cuda/fmha_blackwell/collective/fmha_fusion.hpp" + +namespace cutlass::fmha::collective { + +using namespace cute; + +template < + class Element, + class StrideQ, + class StrideK, + class StrideV, + class CollectiveMmaQK, + class CollectiveMmaPV, + class SmemLayoutQ, + class SmemLayoutK, + class SmemLayoutV, + class TensorStorage, + class PipelineQ, + class PipelineKV, + class Mask, + class TileShape> +struct Sm100FmhaLoadTmaWarpspecialized { + using TileShapeQK = typename CollectiveMmaQK::TileShape; + using TileShapePV = typename CollectiveMmaPV::TileShape; + + struct Arguments { + const Element* ptr_Q; + StrideQ dQ; + const Element* ptr_K; + StrideK dK; + const Element* ptr_V; + StrideV dV; + }; + + 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; + TMA_K tma_load_k; + TMA_V tma_load_v; + }; + + template + static Params to_underlying_arguments( + ProblemShape const& problem_shape, + Arguments const& args, + void* workspace) { + auto ptr_Q = args.ptr_Q; + auto ptr_K = args.ptr_K; + auto ptr_V = args.ptr_V; + auto dQ = args.dQ; + auto dK = args.dK; + auto dV = args.dV; + + using IntProblemShape = + cute::tuple, int>>; + + IntProblemShape problem_shape_qk; + if constexpr (is_variable_length_v>) { + auto cumulative_length_q = get<0>(problem_shape).cumulative_length; + auto cumulative_length_k = get<1>(problem_shape).cumulative_length; + if (cumulative_length_q != nullptr && cumulative_length_k != nullptr) { + get<0>(problem_shape_qk) = get<0>(problem_shape).total_length; + get<1>(problem_shape_qk) = get<1>(problem_shape).total_length; + get<2>(problem_shape_qk) = get<2>(problem_shape); + get<3>(problem_shape_qk) = get<3>(problem_shape); + } + } else { + problem_shape_qk = problem_shape; + } + + auto params_qk = CollectiveMmaQK::to_underlying_arguments( + problem_shape_qk, + typename CollectiveMmaQK::Arguments{ + ptr_Q, + dQ, + ptr_K, + dK, + }, + /*workspace=*/nullptr); + + auto problem_shape_pv = select<0, 2, 1, 3>(problem_shape_qk); + auto params_pv = CollectiveMmaPV::to_underlying_arguments( + problem_shape_pv, + typename CollectiveMmaPV::Arguments{ + ptr_K, + dK, // never used, dummy + ptr_V, + select<1, 0, 2>(dV), + }, + /*workspace=*/nullptr); + + return Params{ + params_qk.tma_load_a, params_qk.tma_load_b, params_pv.tma_load_b}; + } + + 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_in, + 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) { + BlkCoord blk_coord_q = blk_coord_in; + BlkCoord blk_coord_kv = blk_coord_in; + + int mask_tile_count = + Mask{}.get_trip_count(blk_coord_in, 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) + + // compute gQ, sQ + // we load 2*get<0>(blk_coord), and 2*get<0>(blk_coord) + 1 + ThrMMA mma_qk = typename CollectiveMmaQK::TiledMma{}.get_slice(0); + Tensor mQ_qdl_p = + params.tma_load_q.get_tma_tensor(select<0, 2, 3>(problem_shape)); + + int q_offs_0 = 0; + + if constexpr (is_variable_length_v< + tuple_element_t<0, ParamsProblemShape>>) { + auto cumulative_length_q = get<0>(params_problem_shape).cumulative_length; + if (cumulative_length_q != nullptr) { + q_offs_0 = cumulative_length_q[get<2, 1>(blk_coord_q)]; + get<2, 1>(blk_coord_q) = 0; + } + } + + Tensor mQ_qdl = domain_offset( + make_coord(q_offs_0, _0{}, make_coord(_0{}, _0{})), mQ_qdl_p); + + Tensor gQ_qdl = local_tile( + mQ_qdl, TileShapeQK{}, make_coord(_, _, _), Step<_1, X, _1>{}); + Tensor tSgQ_qdl = mma_qk.partition_A(gQ_qdl); + Tensor sQ = + make_tensor(make_smem_ptr(storage.smem_q.data()), SmemLayoutQ{}); + auto [tQgQ_qdl, tQsQ] = tma_partition( + params.tma_load_q, + _0{}, + make_layout(_1{}), + group_modes<0, 3>(sQ), + group_modes<0, 3>(tSgQ_qdl)); + Tensor tQgQ = tQgQ_qdl(_, _, _0{}, get<2>(blk_coord_q)); + + // compute gK, sK + Tensor mK_kdl_p = + params.tma_load_k.get_tma_tensor(select<1, 2, 3>(problem_shape)); + + int kv_offs_0 = 0; + + if constexpr (is_variable_length_v< + tuple_element_t<1, ParamsProblemShape>>) { + auto cumulative_length = get<1>(params_problem_shape).cumulative_length; + if (cumulative_length != nullptr) { + kv_offs_0 = cumulative_length[get<2, 1>(blk_coord_kv)]; + get<2, 1>(blk_coord_kv) = 0; + } + } + + Tensor mK_kdl = domain_offset( + make_coord(kv_offs_0, _0{}, make_coord(_0{}, _0{})), mK_kdl_p); + + Tensor gK_kdl = local_tile( + mK_kdl, TileShapeQK{}, make_coord(_, _, _), Step{}); + Tensor tSgK_kdl = mma_qk.partition_B(gK_kdl); + Tensor sK = + make_tensor(make_smem_ptr(storage.smem_k.data()), SmemLayoutK{}); + auto [tKgK_kdl, tKsK] = tma_partition( + params.tma_load_k, + _0{}, + make_layout(_1{}), + group_modes<0, 3>(sK), + group_modes<0, 3>(tSgK_kdl)); + Tensor tKgK = tKgK_kdl(_, _, _0{}, get<2>(blk_coord_kv)); + + // compute gV, sV + ThrMMA mma_pv = typename CollectiveMmaPV::TiledMma{}.get_slice(0); + Tensor mV_dkl_p = + params.tma_load_v.get_tma_tensor(select<2, 1, 3>(problem_shape)); + + Tensor mV_dkl = domain_offset( + make_coord(_0{}, kv_offs_0, make_coord(_0{}, _0{})), mV_dkl_p); + + Tensor gV_dkl = local_tile( + mV_dkl, TileShapePV{}, make_coord(_, _, _), Step{}); + Tensor tOgV_dkl = mma_pv.partition_B(gV_dkl); + Tensor sV = + make_tensor(make_smem_ptr(storage.smem_v.data()), SmemLayoutV{}); + auto [tVgV_dkl, tVsV] = tma_partition( + params.tma_load_v, + _0{}, + make_layout(_1{}), + group_modes<0, 3>(sV), + group_modes<0, 3>(tOgV_dkl)); + auto tVgV = tVgV_dkl(_, _0{}, _, get<2>(blk_coord_kv)); + + // 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_q); + int q1_index = 2 * get<0>(blk_coord_q) + 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_kv.producer_acquire(pipeline_kv_producer_state); + if (lane_predicate) { + auto tma_barrier = + pipeline_kv.producer_get_barrier(pipeline_kv_producer_state); + copy( + params.tma_load_k.with(*tma_barrier, 0), + tKgK(_, k_index), + tKsK(_, pipeline_kv_producer_state.index())); + } + ++pipeline_kv_producer_state; + + // 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 + pipeline_kv.producer_acquire(pipeline_kv_producer_state); + if (lane_predicate) { + auto tma_barrier = + pipeline_kv.producer_get_barrier(pipeline_kv_producer_state); + copy( + params.tma_load_v.with(*tma_barrier, 0), + tVgV(_, k_index), + tVsV(_, pipeline_kv_producer_state.index())); + } + ++pipeline_kv_producer_state; + k_index += 1; + + // loop: + mask_tile_count -= 1; + for (; mask_tile_count > 0; mask_tile_count -= 1) { + // Ki + pipeline_kv.producer_acquire(pipeline_kv_producer_state); + if (lane_predicate) { + auto tma_barrier = + pipeline_kv.producer_get_barrier(pipeline_kv_producer_state); + copy( + params.tma_load_k.with(*tma_barrier, 0), + tKgK(_, k_index), + tKsK(_, pipeline_kv_producer_state.index())); + } + ++pipeline_kv_producer_state; + + // Vi + pipeline_kv.producer_acquire(pipeline_kv_producer_state); + if (lane_predicate) { + auto tma_barrier = + pipeline_kv.producer_get_barrier(pipeline_kv_producer_state); + copy( + params.tma_load_v.with(*tma_barrier, 0), + tVgV(_, k_index), + tVsV(_, pipeline_kv_producer_state.index())); + } + ++pipeline_kv_producer_state; + k_index += 1; + } + } +}; + +} // namespace cutlass::fmha::collective diff --git a/natten/csrc/include/natten/cuda/fmha_blackwell/common/pow_2.hpp b/natten/csrc/include/natten/cuda/fmha_blackwell/common/pow_2.hpp new file mode 100644 index 00000000..8b83166a --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_blackwell/common/pow_2.hpp @@ -0,0 +1,92 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2026 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/natten/csrc/include/natten/cuda/fmha_blackwell/device/fmha_bwd_sm100.hpp b/natten/csrc/include/natten/cuda/fmha_blackwell/device/fmha_bwd_sm100.hpp new file mode 100644 index 00000000..5240a7bc --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_blackwell/device/fmha_bwd_sm100.hpp @@ -0,0 +1,446 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 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 +#include "cute/tensor.hpp" +#include "cutlass/cutlass.h" +#include "cutlass/kernel_hardware_info.hpp" + +#include "natten/cuda/fmha_blackwell/device/fmha_sm100.hpp" +#include "natten/cuda/fmha_blackwell/kernel/fmha_kernel_bwd_convert.hpp" +#include "natten/cuda/fmha_blackwell/kernel/fmha_kernel_bwd_sum_OdO.hpp" +#include "natten/cuda/fmha_blackwell/kernel/sm100_fmha_bwd_kernel_tma_warpspecialized.hpp" + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::fmha::device { + +//////////////////////////////////////////////////////////////////////////////// +////////////////////////////// CUTLASS 3.x API ///////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +template < + class ProblemShape, + class Element, + class ElementAccumulator, + class TileShape, + // bool IsMla, + class Mask, + bool IsDeterministic = false> +class FmhaBwdSm100 { + private: + template + constexpr static auto to_bwd_shape(T shape) { + // if constexpr (IsMla) { // remove GQA mode + // constexpr int R = decltype(rank(shape))::value; + // auto HB = get(shape); + // auto rest = take<0,R-1>(shape); + // return append(rest, make_shape(size<0>(HB), get<1>(HB))); + // } + // else { + return shape; + //} + } + + template + constexpr static auto to_bwd_stride(T stride) { + // if constexpr (IsMla) { // remove GQA mode + // constexpr int R = decltype(rank(stride))::value; + // auto HB = get(stride); + // auto rest = take<0,R-1>(stride); + // if constexpr (is_same_v(HB))>, _0>) { + // return append(rest, make_stride(get<0,1>(HB), get<1>(HB))); + // } + // else { + // return append(rest, make_stride(get<0,0>(HB), get<1>(HB))); + // } + // } + // else { + return stride; + //} + } + + public: + /// Argument structure: User API + struct Arguments { + // Q K D D_VO HB + ProblemShape problem_shape; + + const Element* ptr_Q; + cute::tuple, int64_t>> + stride_Q; + const Element* ptr_K; + cute::tuple, int64_t>> + stride_K; + const Element* ptr_V; + cute::tuple, int64_t>> + stride_V; + + const Element* ptr_O; + cute::tuple, int64_t>> + stride_O; + const ElementAccumulator* ptr_LSE; + // NATTEN has a different LSE layout + cute::tuple, int64_t>> stride_LSE; + + const Element* ptr_dO; + cute::tuple, int64_t>> + stride_dO; + + Element* ptr_dQ; + cute::tuple, int64_t>> + stride_dQ; + Element* ptr_dK; + cute::tuple, int64_t>> + stride_dK; + Element* ptr_dV; + cute::tuple, int64_t>> + stride_dV; + + ElementAccumulator softmax_scale; + + int* ptr_dq_semaphore = nullptr; + + cutlass::KernelHardwareInfo hw_info; + }; + + using OperationSumOdO = cutlass::fmha::device::FmhaSm100< + cutlass::fmha::kernel:: + FmhaKernelBwdSumOdO>; + using OperationConvert = cutlass::fmha::device::FmhaSm100< + cutlass::fmha::kernel:: + FmhaKernelBwdConvert>; + + using Operation = cutlass::fmha::device::FmhaSm100< + cutlass::fmha::kernel::Sm100FmhaBwdKernelTmaWarpSpecialized< + ProblemShape, + Element, + ElementAccumulator, + TileShape, + Mask, + IsDeterministic>>; + using Kernel = typename Operation::Kernel; + + struct Params { + OperationSumOdO op_sum_OdO; + Operation op; + OperationConvert op_convert; + ElementAccumulator* dQ_acc; + size_t dQ_acc_size; + }; + + private: + Params params_; + + // Scaled LSE and sum OdO are batch packed, not sequence packed, to avoid + // alignment issues with LSE. + static typename OperationSumOdO::Arguments to_sum_OdO_arguments( + Arguments const& args, + ElementAccumulator* sum_OdO = nullptr, + ElementAccumulator* scaled_lse = nullptr) { + using namespace cute; + auto [Q_, K, D, D_VO, HB] = args.problem_shape; + auto [H, B] = HB; + auto [H_R, H_K] = H; + D = cutlass::round_up(D, 8); // Alignment + int Q = cutlass::round_up(static_cast(Q_), 8); // Alignment + auto stride_scalar = make_stride( + _1{}, + make_stride( + make_stride(Q, Q * H_R), + static_cast(Q) * static_cast(H_R * H_K))); + auto log2_e = log2f(expf(1.0f)); + return typename OperationSumOdO::Arguments{ + args.problem_shape, + args.ptr_O, + args.stride_O, + args.ptr_dO, + args.stride_dO, + sum_OdO, + stride_scalar, + args.ptr_LSE, + args.stride_LSE, + scaled_lse, + stride_scalar, + -1.0f, + -log2_e}; + } + + // F32 dQ is batch-packed, not sequence-packed + static typename OperationConvert::Arguments to_convert_arguments( + Arguments const& args, + ElementAccumulator* src = nullptr) { + using namespace cute; + auto [Q_, K, D, D_VO, HB] = args.problem_shape; + auto [H, B] = HB; + auto [H_R, H_K] = H; + D = cutlass::round_up(D, 8); // Alignment + int Q = cutlass::round_up(static_cast(Q_), 8); // Alignment + auto stride_src_dQ = make_stride( + D, + _1{}, + make_stride( + make_stride(D * Q, D * Q * H_R), + B == 1 ? 0 + : static_cast(D * H_R * H_K) * + static_cast(Q))); + return typename OperationConvert::Arguments{ + args.problem_shape, + src, + stride_src_dQ, + nullptr, + args.stride_dK, + nullptr, + args.stride_dV, + args.ptr_dQ, + args.stride_dQ, + nullptr, + args.stride_dK, + nullptr, + args.stride_dV, + args.softmax_scale}; + } + + static typename Operation::Arguments to_bwd_arguments( + Arguments const& args, + ElementAccumulator* sum_OdO = nullptr, + cute::tuple, int64_t>> const& + stride_sum_OdO = {}, + ElementAccumulator* scaled_lse = nullptr, + cute::tuple, int64_t>> const& + stride_scaled_lse = {}, + ElementAccumulator* dQ_acc = nullptr, + cute::tuple< + int, + cute::_1, + cute::tuple, int64_t>> const& stride_dQ = {}) { + return typename Operation::Arguments{ + to_bwd_shape(args.problem_shape), + {args.ptr_Q, + to_bwd_stride(args.stride_Q), + args.ptr_K, + to_bwd_stride(args.stride_K), + args.ptr_V, + to_bwd_stride(args.stride_V), + args.ptr_dO, + to_bwd_stride(args.stride_dO), + scaled_lse, + to_bwd_stride(stride_scaled_lse), + sum_OdO, + to_bwd_stride(stride_sum_OdO), + dQ_acc, + to_bwd_stride(stride_dQ), + args.softmax_scale, + args.ptr_dq_semaphore}, + {args.ptr_dK, + to_bwd_stride(args.stride_dK), + args.ptr_dV, + to_bwd_stride(args.stride_dV)}, + args.hw_info}; + } + + public: + /// Determines whether the GEMM can execute the given problem. + static Status can_implement(Arguments const& args) { + Status status = Status::kSuccess; + + status = OperationSumOdO::can_implement(to_sum_OdO_arguments(args)); + if (status != Status::kSuccess) { + return status; + } + + status = OperationConvert::can_implement(to_convert_arguments(args)); + if (status != Status::kSuccess) { + return status; + } + + status = Operation::can_implement(to_bwd_arguments(args)); + if (status != Status::kSuccess) { + return status; + } + + return status; + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const& args) { + auto [Q_, K, D_, D_VO, HB] = args.problem_shape; + auto [H_, B_] = product_each(HB); + size_t B = static_cast(B_); + size_t H = static_cast(H_); + size_t D = cutlass::round_up(static_cast(D_), 8); // Alignment + size_t Q = cutlass::round_up(static_cast(Q_), 8); // Alignment + size_t workspace_bytes = 0; + // All three intermediary tensors are batch-packed, not sequence packed + // OdO vector + workspace_bytes += B * H * Q * sizeof(ElementAccumulator); + // scaled LSE vector + workspace_bytes += B * H * Q * sizeof(ElementAccumulator); + // FP32 versions of outputs that are churned (start off with Q only) + workspace_bytes += B * H * Q * D * sizeof(ElementAccumulator); + return workspace_bytes; + } + + /// Initializes state from arguments. + Status initialize_split( + Arguments const& args, + void* workspace_dQ, + void* workspace_sum_OdO, + void* workspace_scaled_lse, + cudaStream_t stream = nullptr) { + CUTLASS_TRACE_HOST( + "FmhaDeviceBwd::initialize_split() - workspace_dQ=" + << workspace_dQ << ", workspace_sum_OdO=" << workspace_sum_OdO + << "stream: " << (stream ? "non-null" : "null")); + + auto [Q_, K, D_, D_VO, HB] = args.problem_shape; + auto [H_, B_] = product_each(HB); + size_t B = static_cast(B_); + size_t H = static_cast(H_); + size_t D = cutlass::round_up(static_cast(D_), 8); // Alignment + size_t Q = cutlass::round_up(static_cast(Q_), 8); // Alignment + ElementAccumulator* sum_OdO = + reinterpret_cast(workspace_sum_OdO); + ElementAccumulator* scaled_lse = + reinterpret_cast(workspace_scaled_lse); + ElementAccumulator* dQ_acc = + reinterpret_cast(workspace_dQ); + params_.dQ_acc = dQ_acc; + params_.dQ_acc_size = B * H * Q * D * sizeof(ElementAccumulator); + auto args_sum_OdO = to_sum_OdO_arguments(args, sum_OdO, scaled_lse); + auto args_convert = to_convert_arguments(args, dQ_acc); + params_.op_sum_OdO.initialize(args_sum_OdO, nullptr, stream); + params_.op_convert.initialize(args_convert, nullptr, stream); + auto args_bwd = to_bwd_arguments( + args, + sum_OdO, + args_sum_OdO.stride_sum_OdO, + scaled_lse, + args_sum_OdO.stride_scaled_lse, + dQ_acc, + args_convert.stride_src_dQ); + params_.op.initialize(args_bwd, nullptr, stream); + + return Status::kSuccess; + } + + /// Initializes state from arguments. + Status initialize( + Arguments const& args, + void* workspace = nullptr, + cudaStream_t stream = nullptr) { + CUTLASS_TRACE_HOST( + "FmhaDeviceBwd::initialize() - workspace " + << workspace << ", stream: " << (stream ? "non-null" : "null")); + + auto [Q_, K, D_, D_VO, HB] = args.problem_shape; + auto [H_, B_] = product_each(HB); + size_t B = static_cast(B_); + size_t H = static_cast(H_); + size_t D = cutlass::round_up(static_cast(D_), 8); // Alignment + size_t Q = cutlass::round_up(static_cast(Q_), 8); // Alignment + char* workspace_chr = reinterpret_cast(workspace); + ElementAccumulator* sum_OdO = + reinterpret_cast(workspace_chr); + workspace_chr += B * H * Q * sizeof(ElementAccumulator); + ElementAccumulator* scaled_lse = + reinterpret_cast(workspace_chr); + workspace_chr += B * H * Q * sizeof(ElementAccumulator); + ElementAccumulator* dQ_acc = + reinterpret_cast(workspace_chr); + return initialize_split(args, dQ_acc, sum_OdO, scaled_lse, stream); + } + + /// 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("FmhaDeviceBwd::run()"); + + Status result = Status::kSuccess; + result = params.op_sum_OdO.run(stream); + if (result != Status::kSuccess) { + return result; + } + + auto cuda_result = + cudaMemsetAsync(params.dQ_acc, 0, params.dQ_acc_size, stream); + if (cuda_result != cudaSuccess) { + return Status::kErrorInternal; + } + + result = params.op.run(stream); + if (result != Status::kSuccess) { + return result; + } + + result = params.op_convert.run(stream); + if (result != Status::kSuccess) { + return result; + } + + 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; + } + + /// 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); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::fmha::device + +//////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fmha_blackwell/device/fmha_sm100.hpp b/natten/csrc/include/natten/cuda/fmha_blackwell/device/fmha_sm100.hpp new file mode 100644 index 00000000..4b7b0301 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_blackwell/device/fmha_sm100.hpp @@ -0,0 +1,286 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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 + +#include +// common +#include "cutlass/cutlass.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 FmhaSm100 { + 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_sm100, + 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_sm100, + 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_sm100, + 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) { + 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_sm100; + void* kernel_params[] = {¶ms}; + launch_result = ClusterLauncher::launch( + grid, cluster, block, smem_size, stream, kernel, kernel_params); + } else { + launch_result = Status::kSuccess; + device_kernel_sm100<<>>(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/natten/csrc/include/natten/cuda/fmha_blackwell/fmha_backward.cuh b/natten/csrc/include/natten/cuda/fmha_blackwell/fmha_backward.cuh new file mode 100644 index 00000000..d64f4958 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_blackwell/fmha_backward.cuh @@ -0,0 +1,255 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ + +#pragma once + +#include "cute/tensor.hpp" + +#include "cutlass/cutlass.h" +#include "cutlass/kernel_hardware_info.h" + +#include "natten/cuda/fmha_blackwell/collective/fmha_fusion.hpp" +#include "natten/cuda/fmha_blackwell/device/fmha_bwd_sm100.hpp" + +namespace natten { +namespace cuda { +namespace fmha_blackwell { + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +using namespace cute; +using namespace cutlass::fmha; + +template < + typename Element, + class TileShape, + class Mask, + bool kIsVarlen, + bool kIsDeterministic = false> +struct KernelBackward { + using ElementAccumulator = float; + using VariableLength = cutlass::fmha::collective::VariableLength; + + // Q K D D_VO ((H_R, H_K) B) + using ProblemShapeRegular = + cute::tuple, int>>; + using ProblemShapeVarlen = cute::tuple< + VariableLength, + VariableLength, + int, + int, + cute::tuple, int>>; + using ProblemShapeType = + std::conditional_t; + + using Operation = cutlass::fmha::device::FmhaBwdSm100< + ProblemShapeType, + Element, + ElementAccumulator, + TileShape, + Mask, + kIsDeterministic>; + + using Arguments = typename Operation::Arguments; + + Operation op; + + Arguments initialize( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch, + int seqlen_Q, + int seqlen_KV, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // varlen parameters + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + int* ptr_dq_semaphore, + // init/launch params + int device_id) { + auto problem_shape_regular = cute::make_tuple( + seqlen_Q, + seqlen_KV, + // head dims are checked to be multiples of 8 for fp16/bf16 and 16 for + // fp8, to meet the 128-bit alignment constraint. No action needed here. + dim, + dim, // dim_value -- if different from dim, needs the MLA kernel + cute::make_tuple(make_tuple(heads_q / heads_kv, heads_kv), batch)); + + ProblemShapeType problem_shape_launch; + decltype(problem_shape_regular) problem_shape_memory; + + if constexpr (kIsVarlen) { + problem_shape_memory = problem_shape_regular; + get<4, 1>(problem_shape_memory) = 1; + + get<0>(problem_shape_launch) = VariableLength{ + max_seqlen_Q, + reinterpret_cast(ptr_cumulative_seqlen_Q), + seqlen_Q}; + get<1>(problem_shape_launch) = VariableLength{ + max_seqlen_KV, + reinterpret_cast(ptr_cumulative_seqlen_KV), + seqlen_KV}; + get<2>(problem_shape_launch) = get<2>(problem_shape_regular); + get<3>(problem_shape_launch) = get<3>(problem_shape_regular); + get<4>(problem_shape_launch) = get<4>(problem_shape_regular); + } else { + problem_shape_memory = problem_shape_regular; + problem_shape_launch = problem_shape_regular; + } + + int SQ = size<0>(problem_shape_memory); + int SK = size<1>(problem_shape_memory); + int D = size<2>(problem_shape_memory); + int D_VO = size<3>(problem_shape_memory); + auto HB = get<4, 0>(problem_shape_memory); + auto [H_R, H_K] = HB; + int B = size<4, 1>(problem_shape_memory); + + // heads last profile, with torch's "contiguous layout" + // shape: (batch, seqlen, heads, dim) + // stride: (dim*heads*seqlen, dim*heads, dim, 1) + auto stride_Q = make_stride( + H_R * H_K * D, + _1{}, + make_stride( + make_stride(D, D * H_R), + B == 1 ? 0L + : static_cast(D * H_R * H_K) * + static_cast(SQ))); + auto stride_K = make_stride( + H_K * D, + _1{}, + make_stride( + make_stride(_0{}, D), + B == 1 ? 0L + : static_cast(D * H_K) * static_cast(SK))); + auto stride_V = make_stride( + H_K * D_VO, + _1{}, + make_stride( + make_stride(_0{}, D_VO), + B == 1 + ? 0L + : static_cast(D_VO * H_K) * static_cast(SK))); + auto stride_O = make_stride( + H_R * H_K * D_VO, + _1{}, + make_stride( + make_stride(D_VO, D_VO * H_R), + B == 1 ? 0L + : static_cast(D_VO * H_R * H_K) * + static_cast(SQ))); + auto stride_LSE = make_stride( + H_K * H_R, + make_stride( + make_stride(_1{}, H_R), + B == 1 + ? 0L + : static_cast(SQ) * static_cast(H_R * H_K))); + + cutlass::KernelHardwareInfo hw_info; + hw_info.device_id = device_id; + hw_info.sm_count = + cutlass::KernelHardwareInfo::query_device_multiprocessor_count( + hw_info.device_id); + + Arguments arguments{ + problem_shape_launch, + reinterpret_cast(ptr_Q), + stride_Q, + reinterpret_cast(ptr_K), + stride_K, + reinterpret_cast(ptr_V), + stride_V, + reinterpret_cast(ptr_O), + stride_O, + reinterpret_cast(ptr_LSE), + stride_LSE, + reinterpret_cast(ptr_dO), + stride_O, + reinterpret_cast(ptr_dQ), + stride_Q, + reinterpret_cast(ptr_dK), + stride_K, + reinterpret_cast(ptr_dV), + stride_V, + attn_scale, + ptr_dq_semaphore, + hw_info}; + + return arguments; + } + + size_t get_workspace_size(Arguments const& arguments) { + return Operation::get_workspace_size(arguments); + } + + void run( + Arguments const& arguments, + void* workspace_ptr, + cudaStream_t stream) { + 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; + return; + } + + status = op.initialize(arguments, workspace_ptr, stream); + if (status != cutlass::Status::kSuccess) { + std::cerr + << "Failed to initialize the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(cudaGetLastError()) << std::endl; + return; + } + + // 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; + } + } +}; + +#endif // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +} // namespace fmha_blackwell +} // namespace cuda +} // namespace natten diff --git a/natten/csrc/include/natten/cuda/fmha_blackwell/fmha_forward.cuh b/natten/csrc/include/natten/cuda/fmha_blackwell/fmha_forward.cuh new file mode 100644 index 00000000..99b82411 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_blackwell/fmha_forward.cuh @@ -0,0 +1,276 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ + +#pragma once + +#include "cute/tensor.hpp" + +#include "cutlass/cutlass.h" +#include "cutlass/kernel_hardware_info.h" + +#include "natten/cuda/fmha_blackwell/collective/fmha_fusion.hpp" +#include "natten/cuda/fmha_blackwell/collective/sm100_fmha_fwd_epilogue_tma_warpspecialized.hpp" +#include "natten/cuda/fmha_blackwell/collective/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp" +#include "natten/cuda/fmha_blackwell/device/fmha_sm100.hpp" +#include "natten/cuda/fmha_blackwell/kernel/fmha_tile_scheduler.hpp" +#include "natten/cuda/fmha_blackwell/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp" + +namespace natten { +namespace cuda { +namespace fmha_blackwell { + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +using namespace cute; +using namespace cutlass::fmha::kernel; +using namespace cutlass::fmha::collective; +using namespace cutlass::fmha; + +template < + typename Element, + class TileShape, + bool kIsPersistent, + class Mask, + bool kIsVarlen> +struct KernelForward { + using ElementAccumulatorQK = float; + using ElementAccumulatorPV = float; + using ElementOut = Element; + using VariableLength = cutlass::fmha::collective::VariableLength; + + // Q K D ((H_R, H_K) B) + using ProblemShapeRegular = + cute::tuple, int>>; + using ProblemShapeVarlen = cute::tuple< + VariableLength, + VariableLength, + int, + cute::tuple, int>>; + using ProblemShapeType = + std::conditional_t; + + using StrideQ = cute:: + tuple, int64_t>>; // Q D (H_G + // H_R B) + using StrideK = cute:: + tuple, int64_t>>; // K D (H_G + // H_R B) + using StrideV = StrideK; + using StrideO = StrideQ; + using StrideLSE = + cute::tuple, int64_t>>; // Q (H_G + // H_R B) + + using TileScheduler = std::conditional_t< + kIsPersistent, + cutlass::fmha::kernel::PersistentTileScheduler, + cutlass::fmha::kernel::IndividualTileScheduler>; + + using Mainloop = + cutlass::fmha::collective::Sm100FmhaFwdMainloopTmaWarpspecialized< + Element, + ElementAccumulatorQK, + ElementAccumulatorPV, + TileShape, + StrideQ, + StrideK, + StrideV, + Mask>; + using Operation = cutlass::fmha::device::FmhaSm100< + cutlass::fmha::kernel::Sm100FmhaFwdKernelTmaWarpspecialized< + ProblemShapeType, + Mainloop, + cutlass::fmha::collective::Sm100FmhaFwdEpilogueTmaWarpspecialized< + ElementOut, + ElementAccumulatorPV, + typename Mainloop::TileShapePV, + StrideO, + StrideLSE>, + TileScheduler>>; + + using Arguments = typename Operation::Arguments; + + Operation op; + + Arguments initialize( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch, + int seqlen_Q, + int seqlen_KV, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // varlen parameters + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id) { + auto problem_shape_regular = cute::make_tuple( + seqlen_Q, + seqlen_KV, + // head dims are checked to be multiples of 8 for fp16/bf16 and 16 for + // fp8, to meet the 128-bit alignment constraint. No action needed here. + dim, + cute::make_tuple( + cute::make_tuple(heads_q / heads_kv, heads_kv), batch)); + + ProblemShapeType problem_shape_launch; + decltype(problem_shape_regular) problem_shape_memory; + + if constexpr (kIsVarlen) { + problem_shape_memory = problem_shape_regular; + get<3, 1>(problem_shape_memory) = 1; + + get<0>(problem_shape_launch) = VariableLength{ + max_seqlen_Q, + reinterpret_cast(ptr_cumulative_seqlen_Q), + seqlen_Q}; + get<1>(problem_shape_launch) = VariableLength{ + max_seqlen_KV, + reinterpret_cast(ptr_cumulative_seqlen_KV), + seqlen_KV}; + get<2>(problem_shape_launch) = get<2>(problem_shape_regular); + get<3>(problem_shape_launch) = get<3>(problem_shape_regular); + } else { + problem_shape_memory = problem_shape_regular; + problem_shape_launch = problem_shape_regular; + } + + int SQ = size<0>(problem_shape_memory); + int SK = size<1>(problem_shape_memory); + int D = size<2>(problem_shape_memory); + int H = size<3, 0>(problem_shape_memory); + int H_K = size<3, 0, 1>(problem_shape_memory); + int H_Q = size<3, 0, 0>(problem_shape_memory); + int B = size<3, 1>(problem_shape_memory); + + // heads last profile, with torch's "contiguous layout" + // shape: (batch, seqlen, heads, dim) + // stride: (dim*heads*seqlen, dim*heads, dim, 1) + auto stride_Q = make_stride( + H * D, + _1{}, + make_stride( + make_stride(D, H_Q * D), + static_cast(H * D) * static_cast(SQ))); + auto stride_O = stride_Q; + auto stride_K = make_stride( + H_K * D, + _1{}, + make_stride( + make_stride(_0{}, D), + static_cast(H_K * D) * static_cast(SK))); + auto stride_V = stride_K; + auto stride_LSE = make_stride( + H, + make_stride( + make_stride(_1{}, H_Q), + static_cast(SQ) * static_cast(H))); + + if (kIsVarlen) { + get<2, 1>(stride_Q) = 0L; + get<2, 1>(stride_K) = 0L; + get<2, 1>(stride_V) = 0L; + get<2, 1>(stride_O) = 0L; + get<1, 1>(stride_LSE) = 0L; + } + + cutlass::KernelHardwareInfo hw_info; + hw_info.device_id = device_id; + hw_info.sm_count = + cutlass::KernelHardwareInfo::query_device_multiprocessor_count( + hw_info.device_id); + + Arguments arguments{ + problem_shape_launch, + {reinterpret_cast(ptr_Q), + stride_Q, + reinterpret_cast(ptr_K), + stride_K, + reinterpret_cast(ptr_V), + stride_V, + attn_scale}, + {reinterpret_cast(ptr_O), + stride_O, + reinterpret_cast(ptr_LSE), + stride_LSE}, + hw_info}; + + return arguments; + } + + size_t get_workspace_size(Arguments const& arguments) { + return Operation::get_workspace_size(arguments); + } + + void run( + Arguments const& arguments, + void* workspace_ptr, + cudaStream_t stream) { + 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; + return; + } + + status = op.initialize(arguments, workspace_ptr, stream); + if (status != cutlass::Status::kSuccess) { + std::cerr + << "Failed to initialize the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(cudaGetLastError()) << std::endl; + return; + } + + // 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; + } + +#if 0 + cudaError_t result = cudaDeviceSynchronize(); + if (result != cudaSuccess) { + std::cerr << "Error running the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(result) << std::endl; + return; + } +#endif + } +}; + +#endif // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +} // namespace fmha_blackwell +} // namespace cuda +} // namespace natten diff --git a/natten/csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_causal_tile_scheduler.hpp b/natten/csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_causal_tile_scheduler.hpp new file mode 100644 index 00000000..a7d961ab --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_causal_tile_scheduler.hpp @@ -0,0 +1,218 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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" + +namespace cutlass::fmha::kernel { + +//////////////////////////////////////////////////////////////////////////////// + +// Swizzle Q tile and H tile to improve L2 cache hit rate, +// and launch the longest main loop first to keep most SMs busy. + +struct CausalIndividualTileScheduler { + static constexpr int TileQ = 16; + static constexpr int TileH = 8; + static constexpr int TileSize = TileQ * TileH; + + struct Params { + dim3 grid; + int tile_max_q; + FastDivmod divmod_tile_col; + FastDivmod divmod_tile_size; + FastDivmod divmod_tile_head; + }; + + bool valid_ = true; + Params params; + + CUTLASS_DEVICE + CausalIndividualTileScheduler(Params const& params) : params(params) {} + + template + static Params to_underlying_arguments( + ProblemSize const& problem_size, + KernelHardwareInfo hw_info, + ClusterShape const& cluster_shape, + TileShape const& tile_shape) { + using namespace cute; + + dim3 grid( + size<3, 0>(problem_size), + round_up( + ceil_div(size<0>(problem_size), size<0>(tile_shape)), + size<0>(cluster_shape)), + size<3, 1>(problem_size)); + // gridDim.x must multiple of TileH + const int tile_col_count = grid.x / TileH; + const int tile_max_q = grid.y / TileQ * TileQ; + return Params{grid, tile_max_q, tile_col_count, TileSize, TileH}; + } + + 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; + const int block_idx = blockIdx.y * gridDim.x + blockIdx.x; + + int tile_idx, tile_tail; + params.divmod_tile_size(tile_idx, tile_tail, block_idx); + + int tile_row_idx, tile_col_idx; + params.divmod_tile_col(tile_row_idx, tile_col_idx, tile_idx); + + int row_offset_in_tail, col_offset_in_tail; + params.divmod_tile_head(row_offset_in_tail, col_offset_in_tail, tile_tail); + + const int row_idx = tile_row_idx * TileQ + row_offset_in_tail; + const int col_idx = tile_col_idx * TileH + col_offset_in_tail; + + // last q tile launch first + if (blockIdx.y >= params.tile_max_q) { + return make_coord( + int(gridDim.y - 1 - blockIdx.y), + _0{}, + make_coord(int(blockIdx.x), int(blockIdx.z))); + } + + return make_coord( + int(gridDim.y) - 1 - row_idx, + _0{}, + make_coord(col_idx, int(blockIdx.z))); + } + + CUTLASS_DEVICE + CausalIndividualTileScheduler& operator++() { + valid_ = false; + return *this; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +//////////////////////////////////////////////////////////////////////////////// + +// Launch order: H Q B +struct CausalPersistentTileScheduler { + struct Params { + int num_blocks; + FastDivmod divmod_h; + FastDivmod divmod_m_block; + FastDivmod divmod_b; + + KernelHardwareInfo hw_info; + }; + + int block_idx = 0; + Params params; + + CUTLASS_DEVICE + CausalPersistentTileScheduler(Params const& params) + : block_idx(blockIdx.x), params(params) {} + + template + static Params to_underlying_arguments( + ProblemSize const& problem_size, + KernelHardwareInfo hw_info, + ClusterShape const& cluster_shape, + TileShape const& tile_shape) { + using namespace cute; + // Get SM count if needed, otherwise use user supplied SM count + int sm_count = hw_info.sm_count; + if (sm_count <= 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 = cutlass::round_up( + ceil_div(size<0>(problem_size), size<0>(tile_shape)), + size<0>(cluster_shape)); + int num_blocks = + num_m_blocks * size<3, 0>(problem_size) * size<3, 1>(problem_size); + + return Params{ + num_blocks, + {size<3, 0>(problem_size)}, + {num_m_blocks}, + {size<3, 1>(problem_size)}, + 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, bidh; + params.divmod_h(block_decode, bidh, block_decode); + params.divmod_m_block(block_decode, m_block, block_decode); + params.divmod_b(block_decode, bidb, block_decode); + return make_coord(m_block, _0{}, make_coord(bidh, bidb)); + } + + CUTLASS_DEVICE + CausalPersistentTileScheduler& operator++() { + block_idx += gridDim.x; + return *this; + } +}; +//////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::fmha::kernel diff --git a/natten/csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_kernel_bwd_convert.hpp b/natten/csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_kernel_bwd_convert.hpp new file mode 100644 index 00000000..98e5e6d3 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_kernel_bwd_convert.hpp @@ -0,0 +1,218 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 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/arch/arch.h" +#include "cutlass/cutlass.h" + +#include "natten/cuda/fmha_blackwell/collective/fmha_fusion.hpp" + +namespace cutlass::fmha::kernel { + +using namespace cutlass::fmha::collective; + +using namespace cute; + +template +struct FmhaKernelBwdConvert { + struct Arguments { + ProblemShape problem_shape; + + const ElementAcc* ptr_src_dQ; + tuple, int64_t>> stride_src_dQ; + const ElementAcc* ptr_src_dK; + tuple, int64_t>> stride_src_dK; + const ElementAcc* ptr_src_dV; + tuple, int64_t>> stride_src_dV; + + Element* ptr_dest_dQ; + tuple, int64_t>> stride_dest_dQ; + Element* ptr_dest_dK; + tuple, int64_t>> stride_dest_dK; + Element* ptr_dest_dV; + tuple, int64_t>> stride_dest_dV; + + ElementAcc scale = 1.0; + }; + + using Params = Arguments; + + using ClusterShape = Shape<_1, _1, _1>; + static constexpr int SharedStorageSize = 0; + + static const int MinBlocksPerMultiprocessor = 1; + static const int MaxThreadsPerBlock = 128; + using ArchTag = cutlass::arch::Sm100; + + static const int kBlockSeq = 8; + + 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 const int kNumThreadsD = 16; + static const int kNumThreadsSeq = MaxThreadsPerBlock / kNumThreadsD; + static const int kElementsPerLoad = 4; + + static const int kIterationsSeq = kBlockSeq / kNumThreadsSeq; + + static bool can_implement(Arguments const& args) { + return get<2>(args.problem_shape) % kElementsPerLoad == 0 && + get<3>(args.problem_shape) % kElementsPerLoad == 0; + } + + static dim3 get_grid_shape(Params const& params) { + dim3 grid( + // never put seq in z, long seqs can easily exceed the 64K limit + ceil_div( + std::max( + size<0>(params.problem_shape), size<1>(params.problem_shape)), + kBlockSeq), + size<4, 0>(params.problem_shape), + size<4, 1>(params.problem_shape)); + return grid; + } + + static dim3 get_block_shape() { + dim3 block(kNumThreadsD, kNumThreadsSeq, 1); + return block; + } + + static Params to_underlying_arguments( + Arguments const& args, + void* workspace) { + return args; + } + + template + CUTLASS_DEVICE void copy( + Params const& params, + const ElementAcc* ptr_src, + StrideSrc const& stride_src, + Element* ptr_dest, + StrideDest const& stride_dest, + Count const& count, + int d_dim) { + auto ptr_src_bh = ptr_src + + (static_cast(get<2, 0, 0>(stride_src)) * + static_cast(blockIdx.y)) + + (static_cast(get<2, 1>(stride_src)) * + static_cast(blockIdx.z)); + auto ptr_dest_bh = ptr_dest + + (static_cast(get<2, 0, 0>(stride_dest)) * + static_cast(blockIdx.y)) + + (static_cast(get<2, 1>(stride_dest)) * + static_cast(blockIdx.z)); + + int seqlen = count; + if constexpr (is_variable_length_v) { + int offset = count.cumulative_length[blockIdx.z]; + ptr_dest_bh += static_cast(offset) * + static_cast(get<0>(stride_dest)); + seqlen = count.cumulative_length[blockIdx.z + 1] - offset; + } + + for (int idx_s_t = threadIdx.y; idx_s_t < kBlockSeq; + idx_s_t += kNumThreadsSeq) { + int idx_s = idx_s_t + kBlockSeq * blockIdx.x; + if (idx_s >= seqlen) + continue; + auto ptr_src_bhs = ptr_src_bh + + (static_cast(idx_s) * + static_cast(get<0>(stride_src))); + auto ptr_dest_bhs = ptr_dest_bh + + (static_cast(idx_s) * + static_cast(get<0>(stride_dest))); + + for (int idx_d = threadIdx.x * kElementsPerLoad; idx_d < d_dim; + idx_d += kElementsPerLoad * kNumThreadsD) { + ElementAcc value_src[kElementsPerLoad]; + Element value_dest[kElementsPerLoad]; + + using VecSrc = uint_bit_t * kElementsPerLoad>; + using VecDest = uint_bit_t * kElementsPerLoad>; + *reinterpret_cast(value_src) = + *reinterpret_cast(&ptr_src_bhs[idx_d]); + + for (int v = 0; v < kElementsPerLoad; v++) { + value_dest[v] = static_cast(params.scale * value_src[v]); + } + + *reinterpret_cast(&ptr_dest_bhs[idx_d]) = + *reinterpret_cast(value_dest); + } + } + } + + CUTLASS_DEVICE void operator()(const Params& params, char* smem) { + if (params.ptr_src_dQ != nullptr) { + copy( + params, + params.ptr_src_dQ, + params.stride_src_dQ, + params.ptr_dest_dQ, + params.stride_dest_dQ, + get<0>(params.problem_shape), + get<2>(params.problem_shape)); + } + if (params.ptr_src_dK != nullptr) { + copy( + params, + params.ptr_src_dK, + params.stride_src_dK, + params.ptr_dest_dK, + params.stride_dest_dK, + get<1>(params.problem_shape), + get<2>(params.problem_shape)); + } + if (params.ptr_src_dV != nullptr) { + copy( + params, + params.ptr_src_dV, + params.stride_src_dV, + params.ptr_dest_dV, + params.stride_dest_dV, + get<1>(params.problem_shape), + get<3>(params.problem_shape)); + } + } +}; + +} // namespace cutlass::fmha::kernel diff --git a/natten/csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_kernel_bwd_sum_OdO.hpp b/natten/csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_kernel_bwd_sum_OdO.hpp new file mode 100644 index 00000000..9756bccb --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_kernel_bwd_sum_OdO.hpp @@ -0,0 +1,194 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 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::kernel { + +using namespace cute; + +template +struct FmhaKernelBwdSumOdO { + struct Arguments { + // [Q, K, D, D_VO, [H, B]] + ProblemShape problem_shape; + + const Element* ptr_O; + cute::tuple, int64_t>> + stride_O; + const Element* ptr_dO; + cute::tuple, int64_t>> + stride_dO; + + ElementAcc* ptr_sum_OdO; + cute::tuple, int64_t>> + stride_sum_OdO; + + const ElementAcc* ptr_lse; + // NATTEN has a different LSE layout + cute::tuple, int64_t>> + stride_lse; + + ElementAcc* ptr_scaled_lse; + cute::tuple, int64_t>> + stride_scaled_lse; + + ElementAcc sum_odo_scale = 1.0; + ElementAcc lse_scale = 1.0; + }; + + using Params = Arguments; + + using ClusterShape = Shape<_1, _1, _1>; + static constexpr int SharedStorageSize = 0; + + static const int MinBlocksPerMultiprocessor = 1; + static const int MaxThreadsPerBlock = 128; + 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 const int kBlockQ = 16; + + static const int kNumThreadsD = 8; + static const int kNumThreadsQ = MaxThreadsPerBlock / kNumThreadsD; + static const int kElementsPerLoad = 2; + + static const int kIterationsQ = kBlockQ / kNumThreadsQ; + + static bool can_implement(Arguments const& args) { + // head_dim and head_dim_v + return get<2>(args.problem_shape) % kElementsPerLoad == 0 && + get<3>(args.problem_shape) % kElementsPerLoad == 0; + } + + static dim3 get_grid_shape(Params const& params) { + dim3 grid( + ceil_div(size<0>(params.problem_shape), kBlockQ), + size<4, 0>(params.problem_shape), + size<4, 1>(params.problem_shape)); + return grid; + } + + static dim3 get_block_shape() { + dim3 block(kNumThreadsD, kNumThreadsQ, 1); + return block; + } + + static Params to_underlying_arguments( + Arguments const& args, + void* workspace) { + return args; + } + + CUTLASS_DEVICE void operator()(const Params& params, char* smem) { + auto ptr_O_bh = params.ptr_O + blockIdx.y * get<2, 0, 0>(params.stride_O) + + blockIdx.z * get<2, 1>(params.stride_O); + auto ptr_dO_bh = params.ptr_dO + + blockIdx.y * get<2, 0, 0>(params.stride_dO) + + blockIdx.z * get<2, 1>(params.stride_dO); + auto ptr_sum_OdO_bh = params.ptr_sum_OdO + + blockIdx.y * get<1, 0, 0>(params.stride_sum_OdO) + + blockIdx.z * get<1, 1>(params.stride_sum_OdO); + auto ptr_lse_bh = params.ptr_lse + + blockIdx.y * get<1, 0, 0>(params.stride_lse) + + blockIdx.z * get<1, 1>(params.stride_lse); + auto ptr_scaled_lse_bh = params.ptr_scaled_lse + + blockIdx.y * get<1, 0, 0>(params.stride_scaled_lse) + + blockIdx.z * get<1, 1>(params.stride_scaled_lse); + + auto problem_q = get<0>(params.problem_shape); + int seqlen_q = problem_q; + if constexpr (is_variable_length_v) { + int offset = problem_q.cumulative_length[blockIdx.z]; + ptr_O_bh += offset * get<0>(params.stride_O); + ptr_dO_bh += offset * get<0>(params.stride_dO); + ptr_lse_bh += offset * get<0>(params.stride_lse); + seqlen_q = problem_q.cumulative_length[blockIdx.z + 1] - offset; + } + + CUTLASS_PRAGMA_UNROLL + for (int idx_q_t = threadIdx.y; idx_q_t < kBlockQ; + idx_q_t += kNumThreadsQ) { + int idx_q = idx_q_t + kBlockQ * blockIdx.x; + if (idx_q >= seqlen_q) + continue; + ElementAcc acc = 0; + auto ptr_O_bhq = ptr_O_bh + idx_q * get<0>(params.stride_O); + auto ptr_dO_bhq = ptr_dO_bh + idx_q * get<0>(params.stride_dO); + auto ptr_sum_OdO_bhq = + ptr_sum_OdO_bh + idx_q * get<0>(params.stride_sum_OdO); + auto ptr_lse_bhq = ptr_lse_bh + idx_q * get<0>(params.stride_lse); + auto ptr_scaled_lse_bhq = + ptr_scaled_lse_bh + idx_q * get<0>(params.stride_scaled_lse); + + for (int idx_d = threadIdx.x * kElementsPerLoad; + idx_d < get<3>(params.problem_shape); + idx_d += kElementsPerLoad * kNumThreadsD) { + Element value_O[kElementsPerLoad]; + Element value_dO[kElementsPerLoad]; + + using Vec = uint_bit_t * kElementsPerLoad>; + *reinterpret_cast(value_O) = + *reinterpret_cast(&ptr_O_bhq[idx_d]); + *reinterpret_cast(value_dO) = + *reinterpret_cast(&ptr_dO_bhq[idx_d]); + + for (int v = 0; v < kElementsPerLoad; v++) { + acc += value_O[v] * value_dO[v]; + } + } + + for (int i = 1; i < kNumThreadsD; i *= 2) { + acc += __shfl_xor_sync((uint32_t)-1, acc, i, kNumThreadsD); + } + + if (threadIdx.x == 0) { + *ptr_sum_OdO_bhq = params.sum_odo_scale * acc; + *ptr_scaled_lse_bhq = params.lse_scale * *ptr_lse_bhq; + } + } + } +}; + +} // namespace cutlass::fmha::kernel diff --git a/natten/csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_tile_scheduler.hpp b/natten/csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_tile_scheduler.hpp new file mode 100644 index 00000000..2a644630 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_tile_scheduler.hpp @@ -0,0 +1,176 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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 IndividualTileScheduler { + struct Params { + dim3 grid; + }; + + bool valid_ = true; + + CUTLASS_DEVICE + IndividualTileScheduler(Params const&) {} + + template + static Params to_underlying_arguments( + ProblemSize const& problem_size, + KernelHardwareInfo hw_info, + ClusterShape const& cluster_shape, + TileShape const& tile_shape) { + using namespace cute; + dim3 grid( + round_up( + ceil_div(size<0>(problem_size), size<0>(tile_shape)), + size<0>(cluster_shape)), + size<3, 0>(problem_size), + size<3, 1>(problem_size)); + 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{}, make_coord(blockIdx.y, blockIdx.z)); + } + + CUTLASS_DEVICE + IndividualTileScheduler& operator++() { + valid_ = false; + return *this; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +struct PersistentTileScheduler { + struct Params { + int num_blocks; + FastDivmod divmod_m_block; + FastDivmod divmod_h; + FastDivmod divmod_b; + + KernelHardwareInfo hw_info; + }; + + int block_idx = 0; + Params params; + + CUTLASS_DEVICE + PersistentTileScheduler(Params const& params) + : block_idx(blockIdx.x), params(params) {} + + template + static Params to_underlying_arguments( + ProblemSize const& problem_size, + KernelHardwareInfo hw_info, + ClusterShape const& cluster_shape, + TileShape const& tile_shape) { + using namespace cute; + // Get SM count if needed, otherwise use user supplied SM count + int sm_count = hw_info.sm_count; + if (sm_count <= 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 = cutlass::round_up( + ceil_div(size<0>(problem_size), size<0>(tile_shape)), + size<0>(cluster_shape)); + int num_blocks = + num_m_blocks * size<3, 0>(problem_size) * size<3, 1>(problem_size); + + return Params{ + num_blocks, + {num_m_blocks}, + {size<3, 0>(problem_size)}, + {size<3, 1>(problem_size)}, + 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, bidh; + params.divmod_m_block(block_decode, m_block, block_decode); + params.divmod_b(block_decode, bidb, block_decode); + params.divmod_h(block_decode, bidh, block_decode); + return make_coord(m_block, _0{}, make_coord(bidh, bidb)); + } + + CUTLASS_DEVICE + PersistentTileScheduler& operator++() { + block_idx += gridDim.x; + return *this; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::fmha::kernel diff --git a/natten/csrc/include/natten/cuda/fmha_blackwell/kernel/sm100_fmha_bwd_kernel_tma_warpspecialized.hpp b/natten/csrc/include/natten/cuda/fmha_blackwell/kernel/sm100_fmha_bwd_kernel_tma_warpspecialized.hpp new file mode 100644 index 00000000..62e648e1 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_blackwell/kernel/sm100_fmha_bwd_kernel_tma_warpspecialized.hpp @@ -0,0 +1,2334 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 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. + * + **************************************************************************************************/ + +/* + * Fmha Bwd deterministic mode was ported from pytorch/FBGEMM + * + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include "cutlass/barrier.h" +#include "cutlass/cutlass.h" + +#include "cute/arch/simd_sm100.hpp" +#include "cute/tensor.hpp" + +#include "cutlass/arch/arch.h" +#include "cutlass/arch/memory_sm80.h" +#include "cutlass/gemm/collective/collective_builder.hpp" + +#include "natten/cuda/fmha_blackwell/collective/fmha_common.hpp" + +#include + +namespace cutlass::fmha::kernel { + +using namespace cutlass::fmha::collective; + +using namespace cute; + +template < + class ProblemShape, + class Element, + class ElementAcc, + class TileShape, + class Mask, + bool IsDeterministic> +struct Sm100FmhaBwdKernelTmaWarpSpecialized { + using TileShapeQ = decltype(get<0>(TileShape{})); + static_assert(std::is_same_v, "tile shape K must be 128"); + using TileShapeK = decltype(get<1>(TileShape{})); + static_assert(std::is_same_v, "tile shape K must be 128"); + using TileShapeDQK = decltype(get<2>(TileShape{})); + using TileShapeDVO = decltype(get<2>(TileShape{})); + + using TmemAllocator = cute::TMEM::Allocator1Sm; + struct TmemAllocation { + static constexpr uint32_t kDK = 0; // TileShapeK x TileShapeDQK x acc + static constexpr uint32_t kDV = + kDK + TileShapeDQK{}; // TileShapeK x TileShapeDVO x acc + static constexpr uint32_t kDQ = + kDV + TileShapeDVO{}; // TileShapeQ x TileShapeDQK x acc + static constexpr uint32_t kDP = kDQ; // TileShapeK x TileShapeQ x inp + static constexpr uint32_t kS = kDQ + max(TileShapeQ{}, TileShapeDQK{}); + static constexpr uint32_t kP = kS; + static constexpr uint32_t kTotal = kS + TileShapeQ{}; + }; + + static_assert( + static_cast(TmemAllocation::kTotal) <= + TmemAllocator::Sm100TmemCapacityColumns, + "using too much tmem"); + + enum class WarpRole { + Empty = 0x0, + Load = 0x1, + Mma = 0x2, + Compute = 0x3, + Reduce = 0x4 + }; + + static constexpr unsigned long long kWarpAssignment = 0x12'3333'3333'4444ull; + static constexpr int kNumComputeWarps = 8; + static constexpr int kNumReduceWarps = 4; + CUTLASS_DEVICE WarpRole warp_idx_to_role(int warp_idx) { + return static_cast((kWarpAssignment >> (4 * warp_idx)) & 0xF); + } + + struct RegisterAllocation { + static constexpr int kWarpgroup0 = 160 - 8; + static constexpr int kWarpgroup1 = 128; + static constexpr int kWarpgroup2 = 96; + static constexpr int kReduce = kWarpgroup0; + static constexpr int kCompute = kWarpgroup1; + static constexpr int kMma = kWarpgroup2; + static constexpr int kEmpty = kWarpgroup2; + static constexpr int kLoad = kWarpgroup2; + + static_assert(kWarpgroup0 + 2 * kWarpgroup1 + kWarpgroup2 <= 512); + }; + + using ArchTag = cutlass::arch::Sm100; + + using ClusterShape = Shape<_1, _1, _1>; + using Schedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + + static constexpr int MinBlocksPerMultiprocessor = 1; + static constexpr int kNumWarps = kNumComputeWarps + kNumReduceWarps + 4; + static constexpr int MaxThreadsPerBlock = NumThreadsPerWarp * kNumWarps; + + static constexpr int Alignment = 128 / sizeof_bits_v; + static constexpr int kStages = 2; + + using TensorStrideContiguousK = + Stride, int64_t>>; + using TensorStrideContiguousMN = + Stride<_1, int, Stride, int64_t>>; + using TensorStrideContiguousK_GQA = + Stride, int64_t>>; + using TensorStrideContiguousMN_GQA = + Stride<_1, int, Stride, int64_t>>; + + // compute S + using CollectiveMmaKQ = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, + cutlass::arch::OpClassTensorOp, + Element, + TensorStrideContiguousK_GQA, + Alignment, + Element, + TensorStrideContiguousK, + Alignment, + ElementAcc, + Shape, + ClusterShape, + cutlass::gemm::collective::StageCount, + Schedule>::CollectiveOp; + using TileShapeKQ = typename CollectiveMmaKQ::TileShape; + using TiledMmaKQ = typename CollectiveMmaKQ::TiledMma; + + // compute dP + using CollectiveMmaVDO = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, + cutlass::arch::OpClassTensorOp, + Element, + TensorStrideContiguousK_GQA, + Alignment, + Element, + TensorStrideContiguousK, + Alignment, + ElementAcc, + Shape, + ClusterShape, + cutlass::gemm::collective::StageCount, + Schedule>::CollectiveOp; + using TileShapeVDO = typename CollectiveMmaVDO::TileShape; + using TiledMmaVDO = typename CollectiveMmaVDO::TiledMma; + + // compute dV + using CollectiveMmaPDO = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, + cutlass::arch::OpClassTensorOp, + // needs to match ordering of S calculation + Element, + TensorStrideContiguousK, + Alignment, + Element, + TensorStrideContiguousMN, + Alignment, + ElementAcc, + Shape, + ClusterShape, + cutlass::gemm::collective::StageCount, + Schedule>::CollectiveOp; + using TileShapePDO = typename CollectiveMmaPDO::TileShape; + using TiledMmaPDO = + decltype(to_tiled_mma_sm100_ts(typename CollectiveMmaPDO::TiledMma{})); + + // compute dK + using CollectiveMmaDSQ = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, + cutlass::arch::OpClassTensorOp, + // somewhat arbitrary since we dump to smem, need to agree with the + // next one + Element, + TensorStrideContiguousK, + Alignment, + Element, + TensorStrideContiguousMN, + Alignment, + ElementAcc, + Shape, + ClusterShape, + cutlass::gemm::collective::StageCount, + Schedule>::CollectiveOp; + using TileShapeDSQ = typename CollectiveMmaDSQ::TileShape; + using TiledMmaDSQ = typename CollectiveMmaDSQ::TiledMma; + + // compute dQ + using CollectiveMmaDSK = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, + cutlass::arch::OpClassTensorOp, + // somewhat arbitrary since we dump to smem, need to agree with the + // previous one + Element, + TensorStrideContiguousMN, + Alignment, + Element, + TensorStrideContiguousMN_GQA, + Alignment, + ElementAcc, + Shape, + ClusterShape, + cutlass::gemm::collective::StageCount, + Schedule>::CollectiveOp; + using TileShapeDSK = typename CollectiveMmaDSK::TileShape; + using TiledMmaDSK = typename CollectiveMmaDSK::TiledMma; + + // pipelines are named Pipeline + static constexpr int kStagesComputeSmem = 1; + using PipelineLoadMmaQ = PipelineTmaUmmaAsync<2, ClusterShape>; + using PipelineLoadMmaDO = PipelineTmaUmmaAsync<1, ClusterShape>; + using PipelineLoadComputeLSE = PipelineAsync<1>; + using PipelineLoadComputeSumOdO = PipelineAsync<1>; + using PipelineMmaComputeS = PipelineUmmaAsync<1>; + using PipelineMmaComputeDP = PipelineUmmaAsync<1>; + using PipelineMmaReduceDQ = PipelineUmmaAsync<1>; + using PipelineComputeMmaP = PipelineUmmaConsumerAsync<1>; + using PipelineComputeMmaDS = PipelineUmmaConsumerAsync; + using PipelineMmaComputeDKDV = PipelineUmmaAsync<2>; + static constexpr int kStagesReduceTmaStore = 2; + using PipelineReduceTmaStore = PipelineTmaStore; + + struct PipelineStorage { + alignas(16) typename PipelineLoadMmaQ::SharedStorage load_mma_q; + alignas(16) typename PipelineLoadMmaDO::SharedStorage load_mma_do; + alignas(16) typename PipelineLoadComputeLSE::SharedStorage load_compute_lse; + alignas(16) + typename PipelineLoadComputeSumOdO::SharedStorage load_compute_sum_odo; + alignas(16) typename PipelineMmaComputeS::SharedStorage mma_compute_s; + alignas(16) typename PipelineMmaComputeDP::SharedStorage mma_compute_dp; + alignas(16) typename PipelineMmaReduceDQ::SharedStorage mma_reduce_dq; + alignas(16) typename PipelineComputeMmaP::SharedStorage compute_mma_p; + alignas(16) typename PipelineComputeMmaDS::SharedStorage compute_mma_ds; + alignas(16) typename PipelineMmaComputeDKDV::SharedStorage mma_compute_dkdv; + }; + + template + static CUTE_DEVICE constexpr auto restage( + Layout const& layout, + Stages stages = {}) { + return composition(layout, make_tuple(_, _, _, make_layout(stages))); + } + + using SmemLayoutK = + decltype(restage(typename CollectiveMmaKQ::SmemLayoutA{})); + using SmemLayoutV = + decltype(restage(typename CollectiveMmaVDO::SmemLayoutA{})); + using SmemLayoutQ = + decltype(restage(typename CollectiveMmaKQ::SmemLayoutB{}, _2{})); + using SmemLayoutDO = + decltype(restage(typename CollectiveMmaVDO::SmemLayoutB{}, _1{})); + using SmemLayoutDS = decltype(restage( + typename CollectiveMmaDSK::SmemLayoutA{}, + Int{})); + using SmemLayoutLSE = Layout>; + using SmemLayoutSumOdO = Layout>; + + using SmemLayoutQT = + decltype(restage(typename CollectiveMmaDSQ::SmemLayoutB{}, _2{})); + using SmemLayoutKT = + decltype(restage(typename CollectiveMmaDSK::SmemLayoutB{})); + using SmemLayoutDST = decltype(restage( + typename CollectiveMmaDSQ::SmemLayoutA{}, + Int{})); + using SmemLayoutDOT = + decltype(restage(typename CollectiveMmaPDO::SmemLayoutB{}, _1{})); + + using TileShapeDQ = _32; + using SmemAtomDQ = + decltype(cutlass::gemm::collective::detail::sm100_smem_selector< + cute::UMMA::Major::K, + ElementAcc, + TileShapeQ, + TileShapeDQ>()); + using SmemShapeDQ = + Shape>; + using SmemLayoutDQ = + decltype(tile_to_shape(SmemAtomDQ{}, SmemShapeDQ{}, Step<_2, _1, _3>{})); + + struct TensorStorage { + union { + alignas(2048) cute::array> smem_k; + alignas(2048) cute::array> smem_k_t; + }; + alignas(2048) cute::array> smem_v; + union { + alignas(2048) cute::array> smem_q; + alignas(2048) cute::array> smem_q_t; + }; + union { + alignas(2048) cute::array> smem_do; + alignas( + 2048) cute::array> smem_do_t; + }; + union { + alignas(2048) cute::array> smem_ds; + alignas( + 2048) cute::array> smem_ds_t; + }; + alignas(1024) cute::array> smem_dq; + alignas(16) cute::array> smem_lse; + alignas(16) + cute::array> smem_sum_odo; + }; + + static constexpr int kTransactionsBytesLoadQ = cutlass::bits_to_bytes( + cosize(take<0, 3>(SmemLayoutQ{})) * cute::sizeof_bits_v); + static constexpr int kTransactionsBytesLoadDO = cutlass::bits_to_bytes( + cosize(take<0, 3>(SmemLayoutDO{})) * cute::sizeof_bits_v); + + static constexpr int kTransactionsBytesLoadK = cutlass::bits_to_bytes( + cosize(take<0, 3>(SmemLayoutK{})) * cute::sizeof_bits_v); + static constexpr int kTransactionsBytesLoadV = cutlass::bits_to_bytes( + cosize(take<0, 3>(SmemLayoutV{})) * cute::sizeof_bits_v); + + struct SharedStorage { + TensorStorage tensors; + PipelineStorage pipelines; + uint32_t tmem_base_ptr; + }; + + // this is tight enough that it won't work with sizeof due to padding for + // alignment + static constexpr int SharedStorageSize = + offsetof(SharedStorage, tmem_base_ptr) + sizeof(uint32_t); + static_assert( + SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, + "using too much smem"); + + using TensorStride = TensorStrideContiguousK; // S D (H B) + using TensorStride_GQA = TensorStrideContiguousK_GQA; + using RowTensorStride = + Stride<_1, Stride, int64_t>>; // S (H B) + + struct MainloopArguments { + const Element* ptr_q; + TensorStride stride_q; + const Element* ptr_k; + TensorStride_GQA stride_k; + const Element* ptr_v; + TensorStride_GQA stride_v; + const Element* ptr_do; + TensorStride stride_do; + + const ElementAcc* ptr_lse; + RowTensorStride stride_lse; + + const ElementAcc* ptr_sum_odo; + RowTensorStride stride_sum_odo; + + ElementAcc* ptr_dq_acc; + TensorStride stride_dq_acc; + + ElementAcc softmax_scale = 1.0f / sqrtf(TileShapeDQK{}); + + int* ptr_dq_semaphore = nullptr; + }; + + using TMA_K = typename CollectiveMmaKQ::Params::TMA_A; + using TMA_V = typename CollectiveMmaVDO::Params::TMA_A; + using TMA_Q = typename CollectiveMmaKQ::Params::TMA_B; + using TMA_DO = typename CollectiveMmaVDO::Params::TMA_B; + + using TMA_DQ = decltype(make_tma_copy( + SM90_TMA_REDUCE_ADD{}, + make_tensor( + (const ElementAcc*)nullptr, + make_shape(1, 1, make_shape(make_shape(1, 1), 1)), + TensorStride{}), + SmemLayoutDQ{}(_, _, _0{}))); + + struct MainloopParams { + TMA_K tma_load_k; + TMA_V tma_load_v; + TMA_Q tma_load_q; + TMA_DO tma_load_do; + TMA_DQ tma_red_dq; + }; + + struct EpilogueArguments { + Element* ptr_dk; + TensorStride_GQA stride_dk; + Element* ptr_dv; + TensorStride_GQA stride_dv; + }; + + struct Arguments { + ProblemShape problem_shape; + MainloopArguments mainloop; + EpilogueArguments epilogue; + KernelHardwareInfo hw_info; + }; + + struct Params { + ProblemShape problem_shape; + MainloopArguments mainloop; + MainloopParams mainloop_params; + EpilogueArguments epilogue; + KernelHardwareInfo hw_info; + }; + + template + CUTLASS_DEVICE int compute_expected_turn( + int iter_index, + int block_k, + ProblemShape_ const& problem_shape) { + if constexpr ( + std::is_base_of_v, Mask> || + std::is_base_of_v, Mask>) { + int offset = 0; + if constexpr (std::is_base_of_v< + cutlass::fmha::collective::CausalMask, + Mask>) { + offset = (get<1>(problem_shape) - get<0>(problem_shape)); + } + int k_global_max = cute::ceil_div(get<1>(problem_shape), TileShapeK{}); + int k_max_for_q_block = std::min( + k_global_max, + cute::ceil_div( + (iter_index + 1) * TileShapeQ{} + offset - 1, TileShapeK{})); + int last_k_block = k_max_for_q_block - 1; + return last_k_block - block_k; + } + return block_k; + } + + static bool can_implement(Arguments const& args) { + auto [Q, K, D, D_VO, HB] = args.problem_shape; + auto [H, B] = HB; + auto [H_R, H_K] = H; + if (Q <= 0 || K <= 0 || D <= 0 || D_VO <= 0 || H_R <= 0 || H_K <= 0 || + B <= 0) { + return false; + } + if (D % Alignment != 0 || D_VO % Alignment != 0) { + return false; + } + return true; + } + + static Status initialize_workspace(Arguments const&, void*, cudaStream_t) { + return Status::kSuccess; + } + + static Params to_underlying_arguments(Arguments const& args, void*) { + auto [Q_, K_, D, D_VO, HB] = args.problem_shape; + int Q = Q_; + int K = K_; + + if constexpr (is_variable_length_v) { + Q = Q_.total_length; + } + if constexpr (is_variable_length_v) { + K = K_.total_length; + } + + auto params_kq = CollectiveMmaKQ::to_underlying_arguments( + make_shape(K, Q, D, HB), + typename CollectiveMmaKQ::Arguments{ + args.mainloop.ptr_k, + args.mainloop.stride_k, + args.mainloop.ptr_q, + args.mainloop.stride_q, + }, + /*workspace=*/nullptr); + + auto params_vdo = CollectiveMmaVDO::to_underlying_arguments( + make_shape(K, Q, D_VO, HB), + typename CollectiveMmaVDO::Arguments{ + args.mainloop.ptr_v, + args.mainloop.stride_v, + args.mainloop.ptr_do, + args.mainloop.stride_do, + }, + /*workspace=*/nullptr); + + TMA_DQ tma_red_dq = make_tma_copy( + SM90_TMA_REDUCE_ADD{}, + make_tensor( + args.mainloop.ptr_dq_acc, + make_shape(Q_, D, HB), + args.mainloop.stride_dq_acc), + SmemLayoutDQ{}(_, _, _0{})); + + return Params{ + args.problem_shape, + args.mainloop, + MainloopParams{ + params_kq.tma_load_a, + params_vdo.tma_load_a, + params_kq.tma_load_b, + params_vdo.tma_load_b, + tma_red_dq}, + args.epilogue, + args.hw_info}; + } + + template + static CUTLASS_DEVICE auto quantize(T const& input) { + constexpr int AlignmentS = 4; + auto output = make_tensor(shape(input)); + auto input_vec = recast>(input); + auto output_vec = recast>(output); + + cutlass::NumericArrayConverter epilogue_op; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(input_vec); i++) { + output_vec(i) = epilogue_op(input_vec(i)); + } + + return output; + } + + template + CUTLASS_DEVICE void load( + BlkCoord const& blk_coord, + BlkOffset const& blk_offset, + ProblemShape_ const& problem_shape, + int iter_start, + int iter_end, + int iter_count, + MainloopArguments const& mainloop_args, + MainloopParams const& mainloop_params, + TensorStorage& shared_tensors, + PipelineLoadMmaQ& pipeline_load_mma_q, + typename PipelineLoadMmaQ::PipelineState& + pipeline_load_mma_q_producer_state, + PipelineLoadMmaDO& pipeline_load_mma_do, + typename PipelineLoadMmaDO::PipelineState& + pipeline_load_mma_do_producer_state, + PipelineLoadComputeLSE& pipeline_load_compute_lse, + typename PipelineLoadComputeLSE::PipelineState& + pipeline_load_compute_lse_producer_state, + PipelineLoadComputeSumOdO& pipeline_load_compute_sum_odo, + typename PipelineLoadComputeSumOdO::PipelineState& + pipeline_load_compute_sum_odo_producer_state) { + auto [Q, K, D, D_VO, HB] = problem_shape; + int iter_index = iter_start; + + using X = Underscore; + + uint16_t mcast_mask = 0; + + auto mK_in = + mainloop_params.tma_load_k.get_tma_tensor(make_shape(K, D, HB)); + auto mV_in = + mainloop_params.tma_load_v.get_tma_tensor(make_shape(K, D_VO, HB)); + auto mQ_in = + mainloop_params.tma_load_q.get_tma_tensor(make_shape(Q, D, HB)); + auto mDO_in = + mainloop_params.tma_load_do.get_tma_tensor(make_shape(Q, D_VO, HB)); + + auto mK = domain_offset(select<1, 2, 4>(blk_offset), mK_in); + auto mV = domain_offset(select<1, 3, 4>(blk_offset), mV_in); + auto mQ = domain_offset(select<0, 2, 4>(blk_offset), mQ_in); + auto mDO = domain_offset(select<0, 3, 4>(blk_offset), mDO_in); + + auto gK = + local_tile(mK, TileShapeKQ{}, make_coord(_, _, _), Step<_1, X, _1>{}); + auto gQ = + local_tile(mQ, TileShapeKQ{}, make_coord(_, _, _), Step{}); + auto gV = + local_tile(mV, TileShapeVDO{}, make_coord(_, _, _), Step<_1, X, _1>{}); + auto gDO = + local_tile(mDO, TileShapeVDO{}, make_coord(_, _, _), Step{}); + + ThrMMA cta_mma_kq = TiledMmaKQ{}.get_slice(_0{}); + ThrMMA cta_mma_vdo = TiledMmaVDO{}.get_slice(_0{}); + + auto tSTgK = cta_mma_kq.partition_A(gK); + auto tSTgQ = cta_mma_kq.partition_B(gQ); + auto tDPTgV = cta_mma_vdo.partition_A(gV); + auto tDPTgDO = cta_mma_vdo.partition_B(gDO); + + auto sQ = make_tensor( + make_smem_ptr(shared_tensors.smem_q.begin()), SmemLayoutQ{}); + auto sK = make_tensor( + make_smem_ptr(shared_tensors.smem_k.begin()), SmemLayoutK{}); + auto sV = make_tensor( + make_smem_ptr(shared_tensors.smem_v.begin()), SmemLayoutV{}); + auto sDO = make_tensor( + make_smem_ptr(shared_tensors.smem_do.begin()), SmemLayoutDO{}); + + auto [tKgK_mkl, tKsK] = tma_partition( + mainloop_params.tma_load_k, + _0{}, + make_layout(_1{}), + group_modes<0, 3>(sK), + group_modes<0, 3>(tSTgK)); + auto [tQgQ_mkl, tQsQ] = tma_partition( + mainloop_params.tma_load_q, + _0{}, + make_layout(_1{}), + group_modes<0, 3>(sQ), + group_modes<0, 3>(tSTgQ)); + auto [tVgV_mkl, tVsV] = tma_partition( + mainloop_params.tma_load_v, + _0{}, + make_layout(_1{}), + group_modes<0, 3>(sV), + group_modes<0, 3>(tDPTgV)); + auto [tDOgDO_mkl, tDOsDO] = tma_partition( + mainloop_params.tma_load_do, + _0{}, + make_layout(_1{}), + group_modes<0, 3>(sDO), + group_modes<0, 3>(tDPTgDO)); + + // set up lse and sum_odo + + auto + [blk_coord_q, blk_coord_k, blk_coord_d, blk_coord_dv, blk_coord_batch] = + blk_coord; + + pipeline_load_mma_q.producer_acquire(pipeline_load_mma_q_producer_state); + auto tma_barrier = pipeline_load_mma_q.producer_get_barrier( + pipeline_load_mma_q_producer_state); + + pipeline_load_mma_q.producer_expect_transaction( + pipeline_load_mma_q_producer_state, kTransactionsBytesLoadK); + + // load K + if (cute::elect_one_sync()) { + cute::copy( + mainloop_params.tma_load_k.with(*tma_barrier, mcast_mask), + tKgK_mkl(_, blk_coord_k, _0{}, blk_coord_batch), + tKsK(_, _0{})); + } + + // load Q + if (cute::elect_one_sync()) { + cute::copy( + mainloop_params.tma_load_q.with(*tma_barrier, mcast_mask), + tQgQ_mkl(_, iter_index, _0{}, blk_coord_batch), + tQsQ(_, pipeline_load_mma_q_producer_state.index())); + } + + ++pipeline_load_mma_q_producer_state; + + pipeline_load_compute_lse.producer_acquire( + pipeline_load_compute_lse_producer_state); + + // load LSE + // 32 threads loading 128 values of 32b each + // so 4*32b=128b + + int thread_idx = threadIdx.x % NumThreadsPerWarp; + int smem_idx = + TileShapeQ{} * pipeline_load_compute_lse_producer_state.index() + + thread_idx * 4; + int gmem_idx = TileShapeQ{} * iter_index + thread_idx * 4; + auto mLSE = make_tensor( + mainloop_args.ptr_lse, make_shape(Q, HB), mainloop_args.stride_lse); + for (int i = 0; i < 4; i++) { + cutlass::arch::cp_async_zfill<4>( + shared_tensors.smem_lse.begin() + smem_idx + i, + &mLSE(gmem_idx + i, blk_coord_batch), + gmem_idx + i < Q); + } + + pipeline_load_compute_lse.producer_commit( + pipeline_load_compute_lse_producer_state, + cutlass::arch::cpasync_barrier_arrive); + ++pipeline_load_compute_lse_producer_state; + + pipeline_load_mma_do.producer_acquire(pipeline_load_mma_do_producer_state); + tma_barrier = pipeline_load_mma_do.producer_get_barrier( + pipeline_load_mma_do_producer_state); + + pipeline_load_mma_do.producer_expect_transaction( + pipeline_load_mma_do_producer_state, kTransactionsBytesLoadV); + + // load V + if (cute::elect_one_sync()) { + cute::copy( + mainloop_params.tma_load_v.with(*tma_barrier, mcast_mask), + tVgV_mkl(_, blk_coord_k, _0{}, blk_coord_batch), + tVsV(_, _0{})); + } + + // load dO + if (cute::elect_one_sync()) { + cute::copy( + mainloop_params.tma_load_do.with(*tma_barrier, mcast_mask), + tDOgDO_mkl(_, iter_index, _0{}, blk_coord_batch), + tDOsDO(_, pipeline_load_mma_do_producer_state.index())); + } + + ++pipeline_load_mma_do_producer_state; + + pipeline_load_compute_sum_odo.producer_acquire( + pipeline_load_compute_sum_odo_producer_state); + + // load sum_OdO + smem_idx = + TileShapeQ{} * pipeline_load_compute_sum_odo_producer_state.index() + + thread_idx * 4; + gmem_idx = TileShapeQ{} * iter_index + thread_idx * 4; + auto mSumOdO = make_tensor( + mainloop_args.ptr_sum_odo, + make_shape(Q, HB), + mainloop_args.stride_sum_odo); + for (int i = 0; i < 4; i++) { + cutlass::arch::cp_async_zfill<4>( + shared_tensors.smem_sum_odo.begin() + smem_idx + i, + &mSumOdO(gmem_idx + i, blk_coord_batch), + gmem_idx + i < Q); + } + + pipeline_load_compute_sum_odo.producer_commit( + pipeline_load_compute_sum_odo_producer_state, + cutlass::arch::cpasync_barrier_arrive); + ++pipeline_load_compute_sum_odo_producer_state; + + iter_count -= 1; + iter_index += 1; + + while (iter_count > 0) { + if (iter_index == iter_end) { + iter_index = iter_start; + get<0, 0>(blk_coord_batch) += 1; + } + + pipeline_load_mma_q.producer_acquire(pipeline_load_mma_q_producer_state); + tma_barrier = pipeline_load_mma_q.producer_get_barrier( + pipeline_load_mma_q_producer_state); + + // load Q + if (cute::elect_one_sync()) { + cute::copy( + mainloop_params.tma_load_q.with(*tma_barrier, mcast_mask), + tQgQ_mkl(_, iter_index, _0{}, blk_coord_batch), + tQsQ(_, pipeline_load_mma_q_producer_state.index())); + } + + ++pipeline_load_mma_q_producer_state; + + pipeline_load_compute_lse.producer_acquire( + pipeline_load_compute_lse_producer_state); + + // load LSE + smem_idx = + TileShapeQ{} * pipeline_load_compute_lse_producer_state.index() + + thread_idx * 4; + gmem_idx = TileShapeQ{} * iter_index + thread_idx * 4; + for (int i = 0; i < 4; i++) { + cutlass::arch::cp_async_zfill<4>( + shared_tensors.smem_lse.begin() + smem_idx + i, + &mLSE(gmem_idx + i, blk_coord_batch), + gmem_idx + i < Q); + } + + pipeline_load_compute_lse.producer_commit( + pipeline_load_compute_lse_producer_state, + cutlass::arch::cpasync_barrier_arrive); + ++pipeline_load_compute_lse_producer_state; + + pipeline_load_mma_do.producer_acquire( + pipeline_load_mma_do_producer_state); + tma_barrier = pipeline_load_mma_do.producer_get_barrier( + pipeline_load_mma_do_producer_state); + + // load dO + if (cute::elect_one_sync()) { + cute::copy( + mainloop_params.tma_load_do.with(*tma_barrier, mcast_mask), + tDOgDO_mkl(_, iter_index, _0{}, blk_coord_batch), + tDOsDO(_, pipeline_load_mma_do_producer_state.index())); + } + + ++pipeline_load_mma_do_producer_state; + + pipeline_load_compute_sum_odo.producer_acquire( + pipeline_load_compute_sum_odo_producer_state); + + // load sum_OdO + smem_idx = + TileShapeQ{} * pipeline_load_compute_sum_odo_producer_state.index() + + thread_idx * 4; + gmem_idx = TileShapeQ{} * iter_index + thread_idx * 4; + for (int i = 0; i < 4; i++) { + cutlass::arch::cp_async_zfill<4>( + shared_tensors.smem_sum_odo.begin() + smem_idx + i, + &mSumOdO(gmem_idx + i, blk_coord_batch), + gmem_idx + i < Q); + } + + pipeline_load_compute_sum_odo.producer_commit( + pipeline_load_compute_sum_odo_producer_state, + cutlass::arch::cpasync_barrier_arrive); + ++pipeline_load_compute_sum_odo_producer_state; + + iter_count -= 1; + iter_index += 1; + } + } + + template + CUTLASS_DEVICE void mma( + BlkCoord const& blk_coord, + ProblemShape_ const& problem_shape, + int iter_start, + int iter_end, + int iter_count, + MainloopArguments const& mainloop_args, + TensorStorage& shared_tensors, + PipelineLoadMmaQ& pipeline_load_mma_q, + typename PipelineLoadMmaQ::PipelineState& + pipeline_load_mma_q_consumer_state, + PipelineLoadMmaDO& pipeline_load_mma_do, + typename PipelineLoadMmaDO::PipelineState& + pipeline_load_mma_do_consumer_state, + PipelineMmaComputeS& pipeline_mma_compute_s, + typename PipelineMmaComputeS::PipelineState& + pipeline_mma_compute_s_producer_state, + PipelineMmaComputeDP& pipeline_mma_compute_dp, + typename PipelineMmaComputeDP::PipelineState& + pipeline_mma_compute_dp_producer_state, + PipelineMmaReduceDQ& pipeline_mma_reduce_dq, + typename PipelineMmaReduceDQ::PipelineState& + pipeline_mma_reduce_dq_producer_state, + PipelineComputeMmaP& pipeline_compute_mma_p, + typename PipelineComputeMmaP::PipelineState& + pipeline_compute_mma_p_consumer_state, + PipelineComputeMmaDS& pipeline_compute_mma_ds, + typename PipelineComputeMmaDS::PipelineState& + pipeline_compute_mma_ds_consumer_state, + PipelineMmaComputeDKDV& pipeline_mma_compute_dkdv, + typename PipelineMmaComputeDKDV::PipelineState& + pipeline_mma_compute_dkdv_producer_state) { + auto [Q, K, D, D_VO, HB] = problem_shape; + + auto sQ = make_tensor( + make_smem_ptr(shared_tensors.smem_q.begin()), SmemLayoutQ{}); + auto sK = make_tensor( + make_smem_ptr(shared_tensors.smem_k.begin()), SmemLayoutK{}); + auto sV = make_tensor( + make_smem_ptr(shared_tensors.smem_v.begin()), SmemLayoutV{}); + auto sDO = make_tensor( + make_smem_ptr(shared_tensors.smem_do.begin()), SmemLayoutDO{}); + + auto sQT = make_tensor( + make_smem_ptr(shared_tensors.smem_q_t.begin()), SmemLayoutQT{}); + auto sKT = make_tensor( + make_smem_ptr(shared_tensors.smem_k_t.begin()), SmemLayoutKT{}); + auto sDS = make_tensor( + make_smem_ptr(shared_tensors.smem_ds.begin()), SmemLayoutDS{}); + auto sDST = make_tensor( + make_smem_ptr(shared_tensors.smem_ds_t.begin()), SmemLayoutDST{}); + auto sP = make_tensor( + make_smem_ptr((Element*)nullptr), + typename CollectiveMmaPDO::SmemLayoutA{}); + auto sDOT = make_tensor( + make_smem_ptr(shared_tensors.smem_do_t.begin()), SmemLayoutDOT{}); + + Tensor tSTrK = TiledMmaKQ::make_fragment_A(sK); + Tensor tSTrQ = TiledMmaKQ::make_fragment_B(sQ); + + Tensor tDPTrV = TiledMmaVDO::make_fragment_A(sV); + Tensor tDPTrDO = TiledMmaVDO::make_fragment_B(sDO); + + Tensor tDQrDS = TiledMmaDSK::make_fragment_A(sDS); + Tensor tDQrKT = TiledMmaDSK::make_fragment_B(sKT); + + Tensor tDKrDST = TiledMmaDSQ::make_fragment_A(sDST); + Tensor tDKrQT = TiledMmaDSQ::make_fragment_B(sQT); + + Tensor tDVrP = TiledMmaPDO::make_fragment_A(sP)(_, _, _, _0{}); + tDVrP.data() = TmemAllocation::kP; + Tensor tDVrDOT = TiledMmaPDO::make_fragment_B(sDOT); + + TiledMmaKQ tiled_mma_kq; + TiledMmaVDO tiled_mma_vdo; + TiledMmaDSK tiled_mma_dsk; + TiledMmaDSQ tiled_mma_dsq; + TiledMmaPDO tiled_mma_pdo; + + tiled_mma_dsq.accumulate_ = UMMA::ScaleOut::Zero; + tiled_mma_pdo.accumulate_ = UMMA::ScaleOut::Zero; + + Tensor tSTtST = + partition_fragment_C(tiled_mma_kq, select<0, 1>(TileShapeKQ{})); + tSTtST.data() = TmemAllocation::kS; + + Tensor tDPTtDPT = + partition_fragment_C(tiled_mma_vdo, select<0, 1>(TileShapeVDO{})); + tDPTtDPT.data() = TmemAllocation::kDP; + + Tensor tDQtDQ = + partition_fragment_C(tiled_mma_dsk, select<0, 1>(TileShapeDSK{})); + tDQtDQ.data() = TmemAllocation::kDQ; + + Tensor tDKtDK = + partition_fragment_C(tiled_mma_dsq, select<0, 1>(TileShapeDSQ{})); + tDKtDK.data() = TmemAllocation::kDK; + + Tensor tDVtDV = + partition_fragment_C(tiled_mma_pdo, select<0, 1>(TileShapePDO{})); + tDVtDV.data() = TmemAllocation::kDV; + + auto pipeline_load_mma_q_release_state = pipeline_load_mma_q_consumer_state; + + pipeline_load_mma_q.consumer_wait(pipeline_load_mma_q_consumer_state); + pipeline_mma_compute_s.producer_acquire( + pipeline_mma_compute_s_producer_state); + + // S = Q*K + tiled_mma_kq.accumulate_ = UMMA::ScaleOut::Zero; + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tSTrQ); ++k_block) { + cute::gemm( + tiled_mma_kq, + tSTrK(_, _, k_block, _0{}), + tSTrQ(_, _, k_block, pipeline_load_mma_q_consumer_state.index()), + tSTtST); + tiled_mma_kq.accumulate_ = UMMA::ScaleOut::One; + } + + ++pipeline_load_mma_q_consumer_state; + + pipeline_mma_compute_s.producer_commit( + pipeline_mma_compute_s_producer_state); + ++pipeline_mma_compute_s_producer_state; + + pipeline_load_mma_do.consumer_wait(pipeline_load_mma_do_consumer_state); + + pipeline_mma_compute_dp.producer_acquire( + pipeline_mma_compute_dp_producer_state); + pipeline_mma_reduce_dq.producer_acquire( + pipeline_mma_reduce_dq_producer_state); + + // dP = dO*V + tiled_mma_vdo.accumulate_ = UMMA::ScaleOut::Zero; + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tDPTrV); ++k_block) { + cute::gemm( + tiled_mma_vdo, + tDPTrV(_, _, k_block, _0{}), + tDPTrDO(_, _, k_block, pipeline_load_mma_do_consumer_state.index()), + tDPTtDPT); + tiled_mma_vdo.accumulate_ = UMMA::ScaleOut::One; + } + + pipeline_mma_compute_dp.producer_commit( + pipeline_mma_compute_dp_producer_state); + ++pipeline_mma_compute_dp_producer_state; + + pipeline_compute_mma_p.consumer_wait(pipeline_compute_mma_p_consumer_state); + + // dV = P*dO + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tDVrP); ++k_block) { + cute::gemm( + tiled_mma_pdo, + tDVrP(_, _, k_block), + tDVrDOT(_, _, k_block, pipeline_load_mma_do_consumer_state.index()), + tDVtDV); + tiled_mma_pdo.accumulate_ = UMMA::ScaleOut::One; + } + + pipeline_compute_mma_p.consumer_release( + pipeline_compute_mma_p_consumer_state); + ++pipeline_compute_mma_p_consumer_state; + + pipeline_load_mma_do.consumer_release(pipeline_load_mma_do_consumer_state); + ++pipeline_load_mma_do_consumer_state; + + iter_count -= 1; + + // in tmem, S & P overlap + // and dP and dQ overlap + // so we need to acquire dQ and dP at the same time + while (iter_count > 0) { + pipeline_load_mma_q.consumer_wait(pipeline_load_mma_q_consumer_state); + pipeline_mma_compute_s.producer_acquire( + pipeline_mma_compute_s_producer_state); + + // S = Q*K + tiled_mma_kq.accumulate_ = UMMA::ScaleOut::Zero; + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tSTrQ); ++k_block) { + cute::gemm( + tiled_mma_kq, + tSTrK(_, _, k_block, _0{}), + tSTrQ(_, _, k_block, pipeline_load_mma_q_consumer_state.index()), + tSTtST); + tiled_mma_kq.accumulate_ = UMMA::ScaleOut::One; + } + + ++pipeline_load_mma_q_consumer_state; + + pipeline_mma_compute_s.producer_commit( + pipeline_mma_compute_s_producer_state); + ++pipeline_mma_compute_s_producer_state; + + pipeline_compute_mma_ds.consumer_wait( + pipeline_compute_mma_ds_consumer_state); + + // we need to acquire dP here, because tmem dQ == tmem dP + pipeline_mma_compute_dp.producer_acquire( + pipeline_mma_compute_dp_producer_state); + + // dQ = dS*K + tiled_mma_dsk.accumulate_ = UMMA::ScaleOut::Zero; + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tDQrDS); ++k_block) { + cute::gemm( + tiled_mma_dsk, + tDQrDS( + _, _, k_block, pipeline_compute_mma_ds_consumer_state.index()), + tDQrKT(_, _, k_block, _0{}), + tDQtDQ); + tiled_mma_dsk.accumulate_ = UMMA::ScaleOut::One; + } + + pipeline_mma_reduce_dq.producer_commit( + pipeline_mma_reduce_dq_producer_state); + ++pipeline_mma_reduce_dq_producer_state; + + // dK = dS*Q + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tDKrDST); ++k_block) { + cute::gemm( + tiled_mma_dsq, + tDKrDST( + _, _, k_block, pipeline_compute_mma_ds_consumer_state.index()), + tDKrQT(_, _, k_block, pipeline_load_mma_q_release_state.index()), + tDKtDK); + tiled_mma_dsq.accumulate_ = UMMA::ScaleOut::One; + } + + pipeline_load_mma_q.consumer_release(pipeline_load_mma_q_release_state); + ++pipeline_load_mma_q_release_state; + + pipeline_compute_mma_ds.consumer_release( + pipeline_compute_mma_ds_consumer_state); + ++pipeline_compute_mma_ds_consumer_state; + + // we grab dq here, because in tmem dq == dp + pipeline_mma_reduce_dq.producer_acquire( + pipeline_mma_reduce_dq_producer_state); + + pipeline_load_mma_do.consumer_wait(pipeline_load_mma_do_consumer_state); + + // dP = dO*V + tiled_mma_vdo.accumulate_ = UMMA::ScaleOut::Zero; + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tDPTrV); ++k_block) { + cute::gemm( + tiled_mma_vdo, + tDPTrV(_, _, k_block, _0{}), + tDPTrDO(_, _, k_block, pipeline_load_mma_do_consumer_state.index()), + tDPTtDPT); + tiled_mma_vdo.accumulate_ = UMMA::ScaleOut::One; + } + + pipeline_mma_compute_dp.producer_commit( + pipeline_mma_compute_dp_producer_state); + ++pipeline_mma_compute_dp_producer_state; + + pipeline_compute_mma_p.consumer_wait( + pipeline_compute_mma_p_consumer_state); + + // dV = P*dO + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tDVrP); ++k_block) { + cute::gemm( + tiled_mma_pdo, + tDVrP(_, _, k_block), + tDVrDOT(_, _, k_block, pipeline_load_mma_do_consumer_state.index()), + tDVtDV); + tiled_mma_pdo.accumulate_ = UMMA::ScaleOut::One; + } + + pipeline_compute_mma_p.consumer_release( + pipeline_compute_mma_p_consumer_state); + ++pipeline_compute_mma_p_consumer_state; + + pipeline_load_mma_do.consumer_release( + pipeline_load_mma_do_consumer_state); + ++pipeline_load_mma_do_consumer_state; + + iter_count -= 1; + } + + // signal to the epilogue that dV is ready + pipeline_mma_compute_dkdv.producer_acquire( + pipeline_mma_compute_dkdv_producer_state); + pipeline_mma_compute_dkdv.producer_commit( + pipeline_mma_compute_dkdv_producer_state); + ++pipeline_mma_compute_dkdv_producer_state; + + pipeline_mma_compute_dkdv.producer_acquire( + pipeline_mma_compute_dkdv_producer_state); + + pipeline_compute_mma_ds.consumer_wait( + pipeline_compute_mma_ds_consumer_state); + + // dK = dS*Q + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tDKrDST); ++k_block) { + cute::gemm( + tiled_mma_dsq, + tDKrDST( + _, _, k_block, pipeline_compute_mma_ds_consumer_state.index()), + tDKrQT(_, _, k_block, pipeline_load_mma_q_release_state.index()), + tDKtDK); + tiled_mma_dsq.accumulate_ = UMMA::ScaleOut::One; + } + + // signal to epilgue that dK is ready + pipeline_mma_compute_dkdv.producer_commit( + pipeline_mma_compute_dkdv_producer_state); + ++pipeline_mma_compute_dkdv_producer_state; + + // we've already acquired mma_reduce_dq in the loop + + // dQ = dS*K + tiled_mma_dsk.accumulate_ = UMMA::ScaleOut::Zero; + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tDQrDS); ++k_block) { + cute::gemm( + tiled_mma_dsk, + tDQrDS(_, _, k_block, pipeline_compute_mma_ds_consumer_state.index()), + tDQrKT(_, _, k_block, _0{}), + tDQtDQ); + tiled_mma_dsk.accumulate_ = UMMA::ScaleOut::One; + } + + pipeline_mma_reduce_dq.producer_commit( + pipeline_mma_reduce_dq_producer_state); + ++pipeline_mma_reduce_dq_producer_state; + + pipeline_load_mma_q.consumer_release(pipeline_load_mma_q_release_state); + ++pipeline_load_mma_q_release_state; + + pipeline_compute_mma_ds.consumer_release( + pipeline_compute_mma_ds_consumer_state); + ++pipeline_compute_mma_ds_consumer_state; + } + + template + CUTLASS_DEVICE void store( + TensorG gmem, + TensorR const& regs, + TensorC const& coord, + TensorShape const& tensor_shape) { + Tensor preds = cute::lazy::transform( + coord, [&](auto const& c) { return elem_less(c, tensor_shape); }); + + auto copy_op = make_cotiled_copy( + Copy_Atom, Element>{}, + make_layout( + make_shape(_1{}, Int{})), + regs.layout()); + auto thr_copy = copy_op.get_slice(_0{}); + + Tensor quantized_regs = quantize(regs); + Tensor tCr = thr_copy.partition_S(quantized_regs); + Tensor tCg = thr_copy.partition_D(gmem); + Tensor tPc = thr_copy.partition_D(preds); + + copy_if(copy_op, tPc, tCr, tCg); + } + + template + CUTLASS_DEVICE void epilogue_clear( + BlkCoord const& blk_coord, + BlkOffset const& blk_offset, + ProblemShape_ const& problem_shape, + MainloopArguments const& mainloop_args, + EpilogueArguments const& epilogue_args) { + auto [Q, K, D, D_VO, HB] = problem_shape; + auto + [blk_coord_q, blk_coord_k, blk_coord_d, blk_coord_dv, blk_coord_batch] = + blk_coord; + + auto mDK_in = make_tensor( + make_gmem_ptr(epilogue_args.ptr_dk), + make_shape(K, TileShapeDQK{}, HB), + epilogue_args.stride_dk); + auto mDK = domain_offset(select<1, 2, 4>(blk_offset), mDK_in); + auto gDK = + local_tile(mDK, TileShapeDSQ{}, make_coord(_, _, _), Step<_1, _1, X>{})( + _, _, blk_coord_k, _0{}, blk_coord_batch); + + Tensor cDK = domain_offset( + make_coord(get<1>(blk_coord) * TileShapeK{}, _0{}), + make_identity_tensor(take<0, 2>(TileShapeDSQ{}))); + + auto mDV_in = make_tensor( + make_gmem_ptr(epilogue_args.ptr_dv), + make_shape(K, TileShapeDVO{}, HB), + epilogue_args.stride_dv); + auto mDV = domain_offset(select<1, 3, 4>(blk_offset), mDV_in); + auto gDV = + local_tile(mDV, TileShapePDO{}, make_coord(_, _, _), Step<_1, _1, X>{})( + _, _, blk_coord_k, _0{}, blk_coord_batch); + + Tensor cDV = domain_offset( + make_coord(blk_coord_k * TileShapeK{}, _0{}), + make_identity_tensor(take<0, 2>(TileShapePDO{}))); + + for (int i = threadIdx.x; i < size(gDK); i += blockDim.x) { + if (elem_less(cDK(i), select<1, 2>(problem_shape))) { + gDK(i) = Element(0); + } + } + for (int i = threadIdx.x; i < size(gDV); i += blockDim.x) { + if (elem_less(cDV(i), select<1, 3>(problem_shape))) { + gDV(i) = Element(0); + } + } + } + + template + CUTLASS_DEVICE void epilogue( + BlkCoord const& blk_coord, + BlkOffset const& blk_offset, + ProblemShape_ const& problem_shape, + MainloopArguments const& mainloop_args, + EpilogueArguments const& epilogue_args, + PipelineMmaComputeDKDV& pipeline_mma_compute_dkdv, + typename PipelineMmaComputeDKDV::PipelineState& + pipeline_mma_compute_dkdv_consumer_state) { + auto [Q, K, D, D_VO, HB] = problem_shape; + auto + [blk_coord_q, blk_coord_k, blk_coord_d, blk_coord_dv, blk_coord_batch] = + blk_coord; + + auto load_op = SM100_TMEM_LOAD_32dp32b16x{}; + + auto tDKtDK = + partition_fragment_C(TiledMmaDSQ{}, select<0, 1>(TileShapeDSQ{}))( + make_coord(_, _), _0{}, _0{}); + tDKtDK.data() = TmemAllocation::kDK; + + auto mDK_in = make_tensor( + make_gmem_ptr(epilogue_args.ptr_dk), + make_shape(K, TileShapeDQK{}, HB), + epilogue_args.stride_dk); + auto mDK = domain_offset(select<1, 2, 4>(blk_offset), mDK_in); + auto gDK = + local_tile(mDK, TileShapeDSQ{}, make_coord(_, _, _), Step<_1, _1, X>{})( + _, _, blk_coord_k, _0{}, blk_coord_batch); + + Tensor cDK = domain_offset( + make_coord(get<1>(blk_coord) * TileShapeK{}, _0{}), + make_identity_tensor(take<0, 2>(TileShapeDSQ{}))); + + constexpr int kNumWarpgroups = kNumComputeWarps / 4; + int dp_idx = threadIdx.x % 128; + int wg_idx = (threadIdx.x % (kNumComputeWarps * NumThreadsPerWarp)) / 128; + + auto split_wg = [&](auto const& t) { + if constexpr (decltype(rank(t))::value == 3) { + auto p = t.compose(make_layout(make_shape( + size<0>(t), + size<1>(t), + make_shape( + Int{}, size<2>(t) / Int{})))); + return p(_, _, make_coord(wg_idx, _)); + } else { + auto p = t.compose(make_layout(make_shape( + size<0>(t), + size<1>(t), + size<2>(t), + make_shape( + Int{}, size<3>(t) / Int{})))); + return p(_, _, _, make_coord(wg_idx, _)); + } + }; + + auto tiled_t2r_dk = make_tmem_copy(load_op, tDKtDK); + auto thread_t2r_dk = tiled_t2r_dk.get_slice(dp_idx); + + Tensor tTR_cDK = split_wg(thread_t2r_dk.partition_D(cDK)); + Tensor tTR_gDK = split_wg(thread_t2r_dk.partition_D(gDK)); + Tensor tTR_rDK = make_tensor(shape(tTR_cDK)); + Tensor tTR_tDK = split_wg(thread_t2r_dk.partition_S(tDKtDK)); + + auto tDVtDV = + partition_fragment_C(TiledMmaDSQ{}, select<0, 1>(TileShapeDSQ{}))( + make_coord(_, _), _0{}, _0{}); + tDVtDV.data() = TmemAllocation::kDV; + + auto mDV_in = make_tensor( + make_gmem_ptr(epilogue_args.ptr_dv), + make_shape(K, TileShapeDVO{}, HB), + epilogue_args.stride_dv); + auto mDV = domain_offset(select<1, 3, 4>(blk_offset), mDV_in); + auto gDV = + local_tile(mDV, TileShapePDO{}, make_coord(_, _, _), Step<_1, _1, X>{})( + _, _, blk_coord_k, _0{}, blk_coord_batch); + + Tensor cDV = domain_offset( + make_coord(blk_coord_k * TileShapeK{}, _0{}), + make_identity_tensor(take<0, 2>(TileShapePDO{}))); + + auto tiled_t2r_dv = make_tmem_copy(load_op, tDVtDV); + auto thread_t2r_dv = tiled_t2r_dv.get_slice(dp_idx); + + Tensor tTR_cDV = split_wg(thread_t2r_dv.partition_D(cDV)); + Tensor tTR_gDV = split_wg(thread_t2r_dv.partition_D(gDV)); + Tensor tTR_rDV = make_tensor(shape(tTR_cDV)); + Tensor tTR_tDV = split_wg(thread_t2r_dv.partition_S(tDVtDV)); + + pipeline_mma_compute_dkdv.consumer_wait( + pipeline_mma_compute_dkdv_consumer_state); + + // load tDVtDV + cute::copy(tiled_t2r_dv, tTR_tDV, tTR_rDV); + + // store tDVgDV + store(tTR_gDV, tTR_rDV, tTR_cDV, select<1, 3>(problem_shape)); + + cutlass::arch::fence_view_async_tmem_load(); + pipeline_mma_compute_dkdv.consumer_release( + pipeline_mma_compute_dkdv_consumer_state); + ++pipeline_mma_compute_dkdv_consumer_state; + + pipeline_mma_compute_dkdv.consumer_wait( + pipeline_mma_compute_dkdv_consumer_state); + + // load tDKtDK + cute::copy(tiled_t2r_dk, tTR_tDK, tTR_rDK); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tTR_rDK); i++) { + tTR_rDK(i) = mainloop_args.softmax_scale * tTR_rDK(i); + } + + // store tDKgDK + store(tTR_gDK, tTR_rDK, tTR_cDK, select<1, 2>(problem_shape)); + + cutlass::arch::fence_view_async_tmem_load(); + pipeline_mma_compute_dkdv.consumer_release( + pipeline_mma_compute_dkdv_consumer_state); + ++pipeline_mma_compute_dkdv_consumer_state; + } + + template + CUTLASS_DEVICE void compute( + BlkCoord const& blk_coord, + BlkOffset const& blk_offset, + ProblemShape_ const& problem_shape, + int iter_start, + int iter_end, + int iter_count, + MainloopArguments const& mainloop_args, + EpilogueArguments const& epilogue_args, + TensorStorage& shared_tensors, + PipelineLoadComputeLSE& pipeline_load_compute_lse, + typename PipelineLoadComputeLSE::PipelineState& + pipeline_load_compute_lse_consumer_state, + PipelineLoadComputeSumOdO& pipeline_load_compute_sum_odo, + typename PipelineLoadComputeSumOdO::PipelineState& + pipeline_load_compute_sum_odo_consumer_state, + PipelineMmaComputeS& pipeline_mma_compute_s, + typename PipelineMmaComputeS::PipelineState& + pipeline_mma_compute_s_consumer_state, + PipelineMmaComputeDP& pipeline_mma_compute_dp, + typename PipelineMmaComputeDP::PipelineState& + pipeline_mma_compute_dp_consumer_state, + PipelineComputeMmaP& pipeline_compute_mma_p, + typename PipelineComputeMmaP::PipelineState& + pipeline_compute_mma_p_producer_state, + PipelineComputeMmaDS& pipeline_compute_mma_ds, + typename PipelineComputeMmaDS::PipelineState& + pipeline_compute_mma_ds_producer_state, + PipelineMmaComputeDKDV& pipeline_mma_compute_dkdv, + typename PipelineMmaComputeDKDV::PipelineState& + pipeline_mma_compute_dkdv_consumer_state) { + auto [Q, K, D, D_VO, HB] = problem_shape; + int iter_index = iter_start; + + // in tmem, S & P overlap + // and dP and dQ overlap + + // there are two compute wg's that cooperatively compute softmax + // they are striped by this tmem atom, i.e. wg0 has 16 elems, then wg1 etc + + auto load_op = SM100_TMEM_LOAD_32dp32b16x{}; + auto store_op = []() { + if constexpr (sizeof(Element) == 1) { + return SM100_TMEM_STORE_32dp32b4x{}; + } else { + return SM100_TMEM_STORE_32dp32b8x{}; + } + }(); + + Tensor tSTtST = + partition_fragment_C(TiledMmaKQ{}, select<0, 1>(TileShapeKQ{}))( + make_coord(_, _), _0{}, _0{}); + tSTtST.data() = TmemAllocation::kS; + + Tensor tDPTtDPT = + partition_fragment_C(TiledMmaVDO{}, select<0, 1>(TileShapeVDO{}))( + make_coord(_, _), _0{}, _0{}); + tDPTtDPT.data() = TmemAllocation::kDP; + + Tensor cST = make_identity_tensor(take<0, 2>(TileShapeKQ{})); + Tensor cDPT = make_identity_tensor(take<0, 2>(TileShapeVDO{})); + + constexpr int kNumWarpgroups = kNumComputeWarps / 4; + int dp_idx = threadIdx.x % 128; + int wg_idx = (threadIdx.x % (kNumComputeWarps * NumThreadsPerWarp)) / 128; + auto tiled_t2r = make_tmem_copy(load_op, tSTtST); + auto thread_t2r = tiled_t2r.get_slice(dp_idx); + + auto split_wg = [&](auto const& t) { + if constexpr (decltype(size<1>(t))::value > 1) { + if constexpr (decltype(rank(t))::value == 3) { + auto p = t.compose(make_layout(make_shape( + size<0>(t), + make_shape( + Int{}, size<1>(t) / Int{}), + size<2>(t)))); + return p(_, make_coord(wg_idx, _), _); + } else { + auto p = t.compose(make_layout(make_shape( + size<0>(t), + make_shape( + Int{}, size<1>(t) / Int{}), + size<2>(t), + size<3>(t)))); + return p(_, make_coord(wg_idx, _), _, _); + } + } else { + if constexpr (decltype(rank(t))::value == 3) { + auto p = t.compose(make_layout(make_shape( + size<0>(t), + size<1>(t), + make_shape( + Int{}, size<2>(t) / Int{})))); + return p(_, _, make_coord(wg_idx, _)); + } else { + auto p = t.compose(make_layout(make_shape( + size<0>(t), + size<1>(t), + size<2>(t), + make_shape( + Int{}, size<3>(t) / Int{})))); + return p(_, _, _, make_coord(wg_idx, _)); + } + } + }; + + Tensor tTR_cST_p = thread_t2r.partition_D(cST); + Tensor tTR_cST = split_wg(tTR_cST_p); + Tensor tTR_rST = make_tensor(shape(tTR_cST)); + // Tensor tTR_tST_p = thread_t2r.partition_S(tSTtST); + Tensor tTR_tST = split_wg(thread_t2r.partition_S(tSTtST)); + + Tensor tTR_cDPT_p = thread_t2r.partition_D(cDPT); + Tensor tTR_cDPT = split_wg(tTR_cDPT_p); + Tensor tTR_rDPT = make_tensor(shape(tTR_cDPT)); + Tensor tTR_tDPT = split_wg(thread_t2r.partition_S(tDPTtDPT)); + + Tensor sLSE = make_tensor( + make_smem_ptr(shared_tensors.smem_lse.begin()), SmemLayoutLSE{}); + Tensor sSumOdO = make_tensor( + make_smem_ptr(shared_tensors.smem_sum_odo.begin()), SmemLayoutSumOdO{}); + + auto sP = make_tensor( + make_smem_ptr((Element*)nullptr), + typename CollectiveMmaPDO::SmemLayoutA{}); + + auto tDVrP = TiledMmaPDO::make_fragment_A(sP)(_, _, _, _0{}); + auto tDVcST = TiledMmaPDO{}.get_slice(_0{}).partition_A(cST); + tDVrP.data() = TmemAllocation::kP; + + auto tiled_r2t = make_tmem_copy(store_op, tDVrP); + auto thread_r2t = tiled_r2t.get_slice(dp_idx); + + auto tRT_tP = split_wg(thread_r2t.partition_D(tDVrP)); + auto tRT_cST_p = thread_r2t.partition_S(tDVcST); + auto tRT_cST = split_wg(tRT_cST_p); + + bool is_residual_k = + get<1>(blk_coord) * TileShapeK{} + TileShapeK{} > get<1>(problem_shape); + + CUTLASS_PRAGMA_NO_UNROLL + while (iter_count > 0) { + // wait for S and P + pipeline_mma_compute_s.consumer_wait( + pipeline_mma_compute_s_consumer_state); + pipeline_compute_mma_p.producer_acquire( + pipeline_compute_mma_p_producer_state); + // wait for LSE + pipeline_load_compute_lse.consumer_wait( + pipeline_load_compute_lse_consumer_state); + + auto dispatch_bool = [](bool b, auto fn) { + if (b) { + fn(cute::true_type{}); + } else { + fn(cute::false_type{}); + } + }; + + bool leading_causal_masking = false; + if constexpr (std::is_base_of_v< + cutlass::fmha::collective::CausalMask, + Mask>) { + leading_causal_masking = warp_uniform(iter_index == iter_start); + } else if constexpr (std::is_base_of_v< + cutlass::fmha::collective::CausalMask, + Mask>) { + int offset = get<1>(problem_shape) - get<0>(problem_shape); + int kv_left = get<1>(blk_coord) * TileShapeK{}; + int kv_right = kv_left + TileShapeK{} - 1; + int q_left = iter_index * TileShapeQ{} + offset; + int q_right = q_left + TileShapeQ{} - 1; + + leading_causal_masking = + warp_uniform(!((q_left > kv_right) || (q_right < kv_left))); + } + bool trailing_residual_masking = false; + if constexpr (std::is_base_of_v< + cutlass::fmha::collective::ResidualMaskForBackward, + Mask>) { + trailing_residual_masking = + warp_uniform((iter_index == iter_end - 1) || is_residual_k); + } + + dispatch_bool( + leading_causal_masking || trailing_residual_masking, + [&](auto is_masked_tile) { + // compute P = softmax(S, LSE) + cute::copy(tiled_t2r, tTR_tST, tTR_rST); + + if constexpr (decltype(is_masked_tile)::value) { + Mask{}.apply_mask( + tTR_rST, + [&](int i) { + auto c_transpose = tTR_cST(i); + return make_coord( + get<1>(c_transpose) + iter_index * TileShapeQ{}, + get<0>(c_transpose) + get<1>(blk_coord) * TileShapeK{}); + }, + problem_shape); + } + + ElementAcc log2_e = static_cast(M_LOG2E); + float2 softmax_scale_log2_e; + softmax_scale_log2_e.x = mainloop_args.softmax_scale * log2_e; + softmax_scale_log2_e.y = mainloop_args.softmax_scale * log2_e; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tTR_rST); i += 2) { + float2 acc; + float2 lse; + float2 out; + acc.x = tTR_rST(i); + acc.y = tTR_rST(i + 1); + lse.x = sLSE( + get<1>(tTR_cST(i)), + pipeline_load_compute_lse_consumer_state.index()); + lse.y = sLSE( + get<1>(tTR_cST(i + 1)), + pipeline_load_compute_lse_consumer_state.index()); + cute::fma(out, softmax_scale_log2_e, acc, lse); + tTR_rST(i) = ::exp2f(out.x); + tTR_rST(i + 1) = ::exp2f(out.y); + } + + auto tRT_rST = quantize(tTR_rST); + auto tRT_rST_reshaped = make_tensor(tRT_rST.data(), shape(tRT_cST)); + + cutlass::arch::fence_view_async_tmem_load(); + cutlass::arch::NamedBarrier( + kNumComputeWarps * NumThreadsPerWarp, + cutlass::arch::ReservedNamedBarriers::TransformBarrier) + .arrive_and_wait(); + + cute::copy(tiled_r2t, tRT_rST_reshaped, tRT_tP); + }); + + // notify for P + cutlass::arch::fence_view_async_tmem_store(); + pipeline_compute_mma_p.producer_commit( + pipeline_compute_mma_p_producer_state); + ++pipeline_compute_mma_p_producer_state; + // release S + pipeline_mma_compute_s.consumer_release( + pipeline_mma_compute_s_consumer_state); + ++pipeline_mma_compute_s_consumer_state; + // release LSE + pipeline_load_compute_lse.consumer_release( + pipeline_load_compute_lse_consumer_state); + ++pipeline_load_compute_lse_consumer_state; + + // wait for OdO + pipeline_load_compute_sum_odo.consumer_wait( + pipeline_load_compute_sum_odo_consumer_state); + // wait for dP + pipeline_mma_compute_dp.consumer_wait( + pipeline_mma_compute_dp_consumer_state); + + // wait for dS + // in principle, we could defer waiting for dS, and move in the freeing of + // dP however, that would force us to keep dS in registers longer + pipeline_compute_mma_ds.producer_acquire( + pipeline_compute_mma_ds_producer_state); + + // compute dS = dsoftmax(P, dP, sum_OdO) + cute::copy(tiled_t2r, tTR_tDPT, tTR_rDPT); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tTR_rDPT); i += 2) { + float2 st; + st.x = tTR_rST(i); + st.y = tTR_rST(i + 1); + float2 dpt; + dpt.x = tTR_rDPT(i); + dpt.y = tTR_rDPT(i + 1); + float2 odo; + odo.x = sSumOdO( + get<1>(tTR_cDPT(i)), + pipeline_load_compute_sum_odo_consumer_state.index()); + odo.y = sSumOdO( + get<1>(tTR_cDPT(i + 1)), + pipeline_load_compute_sum_odo_consumer_state.index()); + float2 dif; + // sum odo is negated during preprocess + cute::add(dif, dpt, odo); + float2 out; + cute::mul(out, dif, st); + tTR_rDPT(i) = out.x; + tTR_rDPT(i + 1) = out.y; + } + + auto tTR_rDST = quantize(tTR_rDPT); + + // release dP + cutlass::arch::fence_view_async_tmem_load(); + pipeline_mma_compute_dp.consumer_release( + pipeline_mma_compute_dp_consumer_state); + ++pipeline_mma_compute_dp_consumer_state; + + Tensor sDS = make_tensor( + make_smem_ptr((Element*)shared_tensors.smem_ds.begin()), + SmemLayoutDS{})( + _, _, _, pipeline_compute_mma_ds_producer_state.index()); + + auto thread_layout = make_ordered_layout( + make_shape(_128{}, _128{}), make_stride(_1{}, _0{})); + + auto sDS_pi = as_position_independent_swizzle_tensor(sDS); + auto sDS_pi_slice_p = sDS_pi.compose(thread_layout)(dp_idx, _).compose( + make_layout(shape(tTR_cDPT_p))); + auto sDS_pi_slice = split_wg(sDS_pi_slice_p); + + copy_aligned(tTR_rDST, sDS_pi_slice); + + // notify for dS + cutlass::arch::fence_view_async_shared(); + pipeline_compute_mma_ds.producer_commit( + pipeline_compute_mma_ds_producer_state); + ++pipeline_compute_mma_ds_producer_state; + // release OdO + pipeline_load_compute_sum_odo.consumer_release( + pipeline_load_compute_sum_odo_consumer_state); + ++pipeline_load_compute_sum_odo_consumer_state; + + iter_count -= 1; + iter_index += 1; + if (iter_index == iter_end) { + iter_index = iter_start; + } + } + + epilogue( + blk_coord, + blk_offset, + problem_shape, + mainloop_args, + epilogue_args, + pipeline_mma_compute_dkdv, + pipeline_mma_compute_dkdv_consumer_state); + } + + template + CUTLASS_DEVICE void reduce( + BlkCoord const& blk_coord, + ProblemShape_ const& problem_shape, + int iter_start, + int iter_end, + int iter_count, + MainloopArguments const& mainloop_args, + MainloopParams const& mainloop_params, + TensorStorage& shared_tensors, + PipelineMmaReduceDQ& pipeline_mma_reduce_dq, + typename PipelineMmaReduceDQ::PipelineState& + pipeline_mma_reduce_dq_consumer_state, + PipelineReduceTmaStore& pipeline_reduce_tma_store, + typename PipelineReduceTmaStore::PipelineState& + pipeline_reduce_tma_store_producer_state) { + using X = Underscore; + + auto [Q, K, D, D_VO, HB] = problem_shape; + auto [H, B] = HB; + auto [H_R, H_K] = H; + int iter_index = iter_start; + + auto + [blk_coord_q, blk_coord_k, blk_coord_d, blk_coord_dv, blk_coord_batch] = + blk_coord; + auto [blx_h, blx_b] = blk_coord_batch; + auto [blx_h_r, blx_h_k] = blx_h; + + // must match TileShapeDQ + auto load_op = SM100_TMEM_LOAD_32dp32b32x{}; + + auto tDQtDQ = + partition_fragment_C(TiledMmaDSK{}, select<0, 1>(TileShapeDSK{}))( + make_coord(_, _), _0{}, _0{}); + tDQtDQ.data() = TmemAllocation::kDQ; + + Tensor mDQ = + mainloop_params.tma_red_dq.get_tma_tensor(make_shape(Q, D, HB)); + auto gDQ = + local_tile(mDQ, TileShapeKQ{}, make_coord(_, _, _), Step{})( + _, _, _, _0{}, _); + + Tensor cDQ = make_identity_tensor(take<0, 2>(TileShapeDSK{})); + + Tensor sDQ = make_tensor( + make_smem_ptr(shared_tensors.smem_dq.begin()), SmemLayoutDQ{}); + + int thread_idx = threadIdx.x % (kNumComputeWarps * NumThreadsPerWarp); + auto tiled_t2r = make_tmem_copy(load_op, tDQtDQ); + auto thread_t2r = tiled_t2r.get_slice(thread_idx); + + Tensor tTR_cDQ = thread_t2r.partition_D(cDQ); + Tensor tTR_sDQ = thread_t2r.partition_D(sDQ); + Tensor tTR_tDQ = thread_t2r.partition_S(tDQtDQ); + + auto block_tma = mainloop_params.tma_red_dq.get_slice(_0{}); + + Tensor tDQsDQ = block_tma.partition_S(sDQ); + Tensor tDQgDQ = block_tma.partition_D(gDQ); + + int lane_predicate = + (threadIdx.x % (kNumReduceWarps * NumThreadsPerWarp)) == 0; + using Barrier = cutlass::GenericBarrier; + int* lock_ptr = !IsDeterministic + ? nullptr + : (mainloop_args.ptr_dq_semaphore + blx_b * H_R * H_K + blx_h_k * H_R); + + int expected_turn = 0; + while (iter_count > 0) { + if constexpr (IsDeterministic) { + expected_turn = + compute_expected_turn(iter_index, blk_coord_k, problem_shape); + Barrier::wait_eq( + lock_ptr, + thread_idx, + iter_index * H_R * H_K * B + get<0, 0>(blk_coord_batch), + expected_turn); + } + { + pipeline_mma_reduce_dq.consumer_wait( + pipeline_mma_reduce_dq_consumer_state); + + Tensor tTR_rDQ = make_tensor(shape(tTR_cDQ)); + + // load dQ from tmem to rmem + cute::copy(tiled_t2r, tTR_tDQ, tTR_rDQ); + + cutlass::arch::fence_view_async_tmem_load(); + pipeline_mma_reduce_dq.consumer_release( + pipeline_mma_reduce_dq_consumer_state); + ++pipeline_mma_reduce_dq_consumer_state; + + // we don't have enough smem to dump it all to smem, so we do it in + // stages + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<2>(tTR_cDQ); i++) { + if (lane_predicate) { + pipeline_reduce_tma_store.producer_acquire( + pipeline_reduce_tma_store_producer_state); + } + // wait in all threads for the acquire to complete + cutlass::arch::NamedBarrier( + kNumReduceWarps * NumThreadsPerWarp, + cutlass::arch::ReservedNamedBarriers::TransposeBarrier) + .arrive_and_wait(); + + cute::copy( + tTR_rDQ(_, _, i), + tTR_sDQ( + _, + _, + _0{}, + pipeline_reduce_tma_store_producer_state.index())); + + // wait for the stores to all be visible to the TMA + cutlass::arch::fence_view_async_shared(); + cutlass::arch::NamedBarrier( + kNumReduceWarps * NumThreadsPerWarp, + cutlass::arch::ReservedNamedBarriers::TransposeBarrier) + .arrive_and_wait(); + if (lane_predicate) { + // launch tma store + copy( + mainloop_params.tma_red_dq, + tDQsDQ( + _, + _, + _0{}, + pipeline_reduce_tma_store_producer_state.index()), + tDQgDQ(_, _, i, iter_index, blk_coord_batch)); + pipeline_reduce_tma_store.producer_commit( + pipeline_reduce_tma_store_producer_state); + } + + ++pipeline_reduce_tma_store_producer_state; + } + } + + if constexpr (IsDeterministic) { + // Increment the semaphore flag + Barrier::arrive_inc( + lock_ptr, + thread_idx, + iter_index * H_R * H_K * B + get<0, 0>(blk_coord_batch)); + } + iter_index += 1; + if (iter_index == iter_end) { + iter_index = iter_start; + get<0, 0>(blk_coord_batch) += 1; + } + + iter_count -= 1; + } + } + + CUTLASS_DEVICE void operator()(Params const& params, char* smem) { +#if ( \ + !defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) && \ + !defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) && \ + !defined(CUTLASS_ARCH_MMA_SM103A_ENABLED) && \ + !defined(CUTLASS_ARCH_MMA_SM103F_ENABLED)) + printf( + "ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); +#else + int warp_idx = cutlass::canonical_warp_idx_sync(); + auto role = warp_idx_to_role(warp_idx); + uint32_t lane_predicate = cute::elect_one_sync(); + + if (role == WarpRole::Load && lane_predicate) { + prefetch_tma_descriptor( + params.mainloop_params.tma_load_q.get_tma_descriptor()); + prefetch_tma_descriptor( + params.mainloop_params.tma_load_k.get_tma_descriptor()); + prefetch_tma_descriptor( + params.mainloop_params.tma_load_v.get_tma_descriptor()); + prefetch_tma_descriptor( + params.mainloop_params.tma_load_do.get_tma_descriptor()); + } + + SharedStorage& shared_storage = *reinterpret_cast(smem); + + int initializing_warp = 0; + typename PipelineLoadMmaQ::Params pipeline_load_mma_q_params; + if (role == WarpRole::Load) { + pipeline_load_mma_q_params.role = + PipelineLoadMmaQ::ThreadCategory::Producer; + } + if (role == WarpRole::Mma) { + pipeline_load_mma_q_params.role = + PipelineLoadMmaQ::ThreadCategory::Consumer; + } + pipeline_load_mma_q_params.is_leader = + lane_predicate && (role == WarpRole::Load); + // Also loads K in the first iteration + pipeline_load_mma_q_params.transaction_bytes = kTransactionsBytesLoadQ; + pipeline_load_mma_q_params.initializing_warp = initializing_warp++; + PipelineLoadMmaQ pipeline_load_mma_q( + shared_storage.pipelines.load_mma_q, + pipeline_load_mma_q_params, + ClusterShape{}, + /*barrier init*/ cute::true_type{}, + /*mask calc*/ cute::false_type{}); + + typename PipelineLoadMmaDO::Params pipeline_load_mma_do_params; + if (role == WarpRole::Load) { + pipeline_load_mma_do_params.role = + PipelineLoadMmaDO::ThreadCategory::Producer; + } + if (role == WarpRole::Mma) { + pipeline_load_mma_do_params.role = + PipelineLoadMmaDO::ThreadCategory::Consumer; + } + pipeline_load_mma_do_params.is_leader = + lane_predicate && (role == WarpRole::Load); + // Also loads V in the first iteration + pipeline_load_mma_do_params.transaction_bytes = kTransactionsBytesLoadDO; + pipeline_load_mma_do_params.initializing_warp = initializing_warp++; + PipelineLoadMmaDO pipeline_load_mma_do( + shared_storage.pipelines.load_mma_do, + pipeline_load_mma_do_params, + ClusterShape{}, + /*barrier init*/ cute::true_type{}, + /*mask calc*/ cute::false_type{}); + + typename PipelineLoadComputeLSE::Params pipeline_load_compute_lse_params; + if (role == WarpRole::Load) { + pipeline_load_compute_lse_params.role = + PipelineLoadComputeLSE::ThreadCategory::Producer; + } + if (role == WarpRole::Compute) { + pipeline_load_compute_lse_params.role = + PipelineLoadComputeLSE::ThreadCategory::Consumer; + } + pipeline_load_compute_lse_params.producer_arv_count = NumThreadsPerWarp; + pipeline_load_compute_lse_params.consumer_arv_count = + kNumComputeWarps * NumThreadsPerWarp; + pipeline_load_compute_lse_params.initializing_warp = initializing_warp++; + PipelineLoadComputeLSE pipeline_load_compute_lse( + shared_storage.pipelines.load_compute_lse, + pipeline_load_compute_lse_params, + /*barrier init*/ cute::true_type{}); + + typename PipelineLoadComputeSumOdO::Params + pipeline_load_compute_sum_odo_params; + if (role == WarpRole::Load) { + pipeline_load_compute_sum_odo_params.role = + PipelineLoadComputeSumOdO::ThreadCategory::Producer; + } + if (role == WarpRole::Compute) { + pipeline_load_compute_sum_odo_params.role = + PipelineLoadComputeSumOdO::ThreadCategory::Consumer; + } + pipeline_load_compute_sum_odo_params.producer_arv_count = NumThreadsPerWarp; + pipeline_load_compute_sum_odo_params.consumer_arv_count = + kNumComputeWarps * NumThreadsPerWarp; + pipeline_load_compute_sum_odo_params.initializing_warp = + initializing_warp++; + PipelineLoadComputeSumOdO pipeline_load_compute_sum_odo( + shared_storage.pipelines.load_compute_sum_odo, + pipeline_load_compute_sum_odo_params, + /*barrier init*/ cute::true_type{}); + + typename PipelineMmaComputeS::Params pipeline_mma_compute_s_params; + if (role == WarpRole::Mma) { + pipeline_mma_compute_s_params.role = + PipelineMmaComputeS::ThreadCategory::Producer; + } + if (role == WarpRole::Compute) { + pipeline_mma_compute_s_params.role = + PipelineMmaComputeS::ThreadCategory::Consumer; + } + pipeline_mma_compute_s_params.consumer_arv_count = + kNumComputeWarps * cutlass::NumThreadsPerWarp; + pipeline_mma_compute_s_params.initializing_warp = initializing_warp++; + PipelineMmaComputeS pipeline_mma_compute_s( + shared_storage.pipelines.mma_compute_s, + pipeline_mma_compute_s_params, + ClusterShape{}, + /*barrier init*/ cute::true_type{}, + /*mask calc*/ cute::false_type{}); + + typename PipelineMmaComputeDP::Params pipeline_mma_compute_dp_params; + if (role == WarpRole::Mma) { + pipeline_mma_compute_dp_params.role = + PipelineMmaComputeDP::ThreadCategory::Producer; + } + if (role == WarpRole::Compute) { + pipeline_mma_compute_dp_params.role = + PipelineMmaComputeDP::ThreadCategory::Consumer; + } + pipeline_mma_compute_dp_params.consumer_arv_count = + kNumComputeWarps * cutlass::NumThreadsPerWarp; + pipeline_mma_compute_dp_params.initializing_warp = initializing_warp++; + PipelineMmaComputeDP pipeline_mma_compute_dp( + shared_storage.pipelines.mma_compute_dp, + pipeline_mma_compute_dp_params, + ClusterShape{}, + /*barrier init*/ cute::true_type{}, + /*mask calc*/ cute::false_type{}); + + typename PipelineMmaReduceDQ::Params pipeline_mma_reduce_dq_params; + if (role == WarpRole::Mma) { + pipeline_mma_reduce_dq_params.role = + PipelineMmaReduceDQ::ThreadCategory::Producer; + } + if (role == WarpRole::Reduce) { + pipeline_mma_reduce_dq_params.role = + PipelineMmaReduceDQ::ThreadCategory::Consumer; + } + pipeline_mma_reduce_dq_params.consumer_arv_count = + kNumReduceWarps * cutlass::NumThreadsPerWarp; + pipeline_mma_reduce_dq_params.initializing_warp = initializing_warp++; + PipelineMmaReduceDQ pipeline_mma_reduce_dq( + shared_storage.pipelines.mma_reduce_dq, + pipeline_mma_reduce_dq_params, + ClusterShape{}, + /*barrier init*/ cute::true_type{}, + /*mask calc*/ cute::false_type{}); + + typename PipelineComputeMmaP::Params pipeline_compute_mma_p_params; + if (role == WarpRole::Mma) { + pipeline_compute_mma_p_params.role = + PipelineComputeMmaP::ThreadCategory::Consumer; + } + if (role == WarpRole::Compute) { + pipeline_compute_mma_p_params.role = + PipelineComputeMmaP::ThreadCategory::Producer; + } + pipeline_compute_mma_p_params.producer_arv_count = + kNumComputeWarps * cutlass::NumThreadsPerWarp; + pipeline_compute_mma_p_params.consumer_arv_count = 1; + pipeline_compute_mma_p_params.initializing_warp = initializing_warp++; + PipelineComputeMmaP pipeline_compute_mma_p( + shared_storage.pipelines.compute_mma_p, + pipeline_compute_mma_p_params, + ClusterShape{}, + /*barrier init*/ cute::true_type{}, + /*mask calc*/ cute::false_type{}); + + typename PipelineComputeMmaDS::Params pipeline_compute_mma_ds_params; + if (role == WarpRole::Mma) { + pipeline_compute_mma_ds_params.role = + PipelineComputeMmaDS::ThreadCategory::Consumer; + } + if (role == WarpRole::Compute) { + pipeline_compute_mma_ds_params.role = + PipelineComputeMmaDS::ThreadCategory::Producer; + } + pipeline_compute_mma_ds_params.producer_arv_count = + kNumComputeWarps * cutlass::NumThreadsPerWarp; + pipeline_compute_mma_ds_params.consumer_arv_count = 1; + pipeline_compute_mma_ds_params.initializing_warp = initializing_warp++; + PipelineComputeMmaDS pipeline_compute_mma_ds( + shared_storage.pipelines.compute_mma_ds, + pipeline_compute_mma_ds_params, + ClusterShape{}, + /*barrier init*/ cute::true_type{}, + /*mask calc*/ cute::false_type{}); + + typename PipelineMmaComputeDKDV::Params pipeline_mma_compute_dkdv_params; + if (role == WarpRole::Mma) { + pipeline_mma_compute_dkdv_params.role = + PipelineMmaComputeDKDV::ThreadCategory::Producer; + } + if (role == WarpRole::Compute) { + pipeline_mma_compute_dkdv_params.role = + PipelineMmaComputeDKDV::ThreadCategory::Consumer; + } + pipeline_mma_compute_dkdv_params.consumer_arv_count = + kNumComputeWarps * cutlass::NumThreadsPerWarp; + pipeline_mma_compute_dkdv_params.initializing_warp = initializing_warp++; + PipelineMmaComputeDKDV pipeline_mma_compute_dkdv( + shared_storage.pipelines.mma_compute_dkdv, + pipeline_mma_compute_dkdv_params, + ClusterShape{}, + /*barrier init*/ cute::true_type{}, + /*mask calc*/ cute::false_type{}); + PipelineReduceTmaStore pipeline_reduce_tma_store; + + TmemAllocator tmem_allocator; + + pipeline_init_arrive_relaxed(size(ClusterShape{})); + + pipeline_load_mma_q.init_masks(ClusterShape{}); + pipeline_load_mma_do.init_masks(ClusterShape{}); + pipeline_mma_compute_s.init_masks(ClusterShape{}); + pipeline_mma_compute_dp.init_masks(ClusterShape{}); + pipeline_mma_reduce_dq.init_masks(ClusterShape{}); + pipeline_compute_mma_p.init_masks(ClusterShape{}); + pipeline_compute_mma_ds.init_masks(ClusterShape{}); + pipeline_mma_compute_dkdv.init_masks(ClusterShape{}); + + typename decltype(pipeline_load_mma_q)::PipelineState + pipeline_load_mma_q_consumer_state; + typename decltype(pipeline_load_mma_do)::PipelineState + pipeline_load_mma_do_consumer_state; + typename decltype(pipeline_load_compute_lse)::PipelineState + pipeline_load_compute_lse_consumer_state; + typename decltype(pipeline_load_compute_sum_odo)::PipelineState + pipeline_load_compute_sum_odo_consumer_state; + typename decltype(pipeline_mma_compute_s)::PipelineState + pipeline_mma_compute_s_consumer_state; + typename decltype(pipeline_mma_compute_dp)::PipelineState + pipeline_mma_compute_dp_consumer_state; + typename decltype(pipeline_mma_reduce_dq)::PipelineState + pipeline_mma_reduce_dq_consumer_state; + typename decltype(pipeline_compute_mma_p)::PipelineState + pipeline_compute_mma_p_consumer_state; + typename decltype(pipeline_compute_mma_ds)::PipelineState + pipeline_compute_mma_ds_consumer_state; + typename decltype(pipeline_mma_compute_dkdv)::PipelineState + pipeline_mma_compute_dkdv_consumer_state; + + auto pipeline_load_mma_q_producer_state = + make_producer_start_state(); + auto pipeline_load_mma_do_producer_state = + make_producer_start_state(); + auto pipeline_load_compute_lse_producer_state = + make_producer_start_state(); + auto pipeline_load_compute_sum_odo_producer_state = + make_producer_start_state(); + auto pipeline_mma_compute_s_producer_state = + make_producer_start_state(); + auto pipeline_mma_compute_dp_producer_state = + make_producer_start_state(); + auto pipeline_mma_reduce_dq_producer_state = + make_producer_start_state(); + auto pipeline_compute_mma_p_producer_state = + make_producer_start_state(); + auto pipeline_compute_mma_ds_producer_state = + make_producer_start_state(); + auto pipeline_mma_compute_dkdv_producer_state = + make_producer_start_state(); + auto pipeline_reduce_tma_store_producer_state = + make_producer_start_state(); + + pipeline_init_wait(size(ClusterShape{})); + + auto blk_coord = make_coord( + _0{}, + blockIdx.x, + _0{}, + _0{}, + make_coord(make_coord(0, blockIdx.y), blockIdx.z)); + auto [problem_shape, blk_offset] = + apply_variable_length_offset(params.problem_shape, blk_coord); + int iter_end = ceil_div(get<0>(problem_shape), TileShapeQ{}); + int iter_start = 0; + if constexpr (std::is_base_of_v< + cutlass::fmha::collective::CausalMask, + Mask>) { + iter_start = (get<1>(blk_coord) * TileShapeK{}) / TileShapeQ{}; + } else if constexpr (std::is_base_of_v< + cutlass::fmha::collective::CausalMask, + Mask>) { + int offset = get<1>(problem_shape) - get<0>(problem_shape); + iter_start = max( + 0, + (int(get<1>(blk_coord) * TileShapeK{}) - offset) / (int)TileShapeQ{}); + } + if (get<1>(blk_coord) * TileShapeK{} >= get<1>(problem_shape)) { + return; + } + int iter_count = (iter_end - iter_start) * get<4, 0, 0>(problem_shape); + + if (iter_count <= 0) { + epilogue_clear( + blk_coord, + blk_offset, + problem_shape, + params.mainloop, + params.epilogue); + return; + } + + if (role == WarpRole::Load) { + warpgroup_reg_set(); + + load( + blk_coord, + blk_offset, + problem_shape, + iter_start, + iter_end, + iter_count, + params.mainloop, + params.mainloop_params, + shared_storage.tensors, + pipeline_load_mma_q, + pipeline_load_mma_q_producer_state, + pipeline_load_mma_do, + pipeline_load_mma_do_producer_state, + pipeline_load_compute_lse, + pipeline_load_compute_lse_producer_state, + pipeline_load_compute_sum_odo, + pipeline_load_compute_sum_odo_producer_state); + + } else if (role == WarpRole::Mma) { + warpgroup_reg_set(); + + tmem_allocator.allocate( + TmemAllocator::Sm100TmemCapacityColumns, + &shared_storage.tmem_base_ptr); + __syncwarp(); + + mma(blk_coord, + problem_shape, + iter_start, + iter_end, + iter_count, + params.mainloop, + shared_storage.tensors, + pipeline_load_mma_q, + pipeline_load_mma_q_consumer_state, + pipeline_load_mma_do, + pipeline_load_mma_do_consumer_state, + pipeline_mma_compute_s, + pipeline_mma_compute_s_producer_state, + pipeline_mma_compute_dp, + pipeline_mma_compute_dp_producer_state, + pipeline_mma_reduce_dq, + pipeline_mma_reduce_dq_producer_state, + pipeline_compute_mma_p, + pipeline_compute_mma_p_consumer_state, + pipeline_compute_mma_ds, + pipeline_compute_mma_ds_consumer_state, + pipeline_mma_compute_dkdv, + pipeline_mma_compute_dkdv_producer_state); + + } else if (role == WarpRole::Compute) { + warpgroup_reg_set(); + + compute( + blk_coord, + blk_offset, + problem_shape, + iter_start, + iter_end, + iter_count, + params.mainloop, + params.epilogue, + shared_storage.tensors, + pipeline_load_compute_lse, + pipeline_load_compute_lse_consumer_state, + pipeline_load_compute_sum_odo, + pipeline_load_compute_sum_odo_consumer_state, + pipeline_mma_compute_s, + pipeline_mma_compute_s_consumer_state, + pipeline_mma_compute_dp, + pipeline_mma_compute_dp_consumer_state, + pipeline_compute_mma_p, + pipeline_compute_mma_p_producer_state, + pipeline_compute_mma_ds, + pipeline_compute_mma_ds_producer_state, + pipeline_mma_compute_dkdv, + pipeline_mma_compute_dkdv_consumer_state); + + cutlass::arch::NamedBarrier( + kNumComputeWarps * NumThreadsPerWarp, + cutlass::arch::ReservedNamedBarriers::EpilogueBarrier) + .arrive_and_wait(); + + if (warp_idx % kNumComputeWarps == 0) { + uint32_t free_stage_ptr = shared_storage.tmem_base_ptr; + tmem_allocator.free( + free_stage_ptr, TmemAllocator::Sm100TmemCapacityColumns); + } + + } else if (role == WarpRole::Reduce) { + warpgroup_reg_set(); + + reduce( + blk_coord, + problem_shape, + iter_start, + iter_end, + iter_count, + params.mainloop, + params.mainloop_params, + shared_storage.tensors, + pipeline_mma_reduce_dq, + pipeline_mma_reduce_dq_consumer_state, + pipeline_reduce_tma_store, + pipeline_reduce_tma_store_producer_state); + + pipeline_reduce_tma_store.producer_tail( + pipeline_reduce_tma_store_producer_state); + } else { + warpgroup_reg_set(); + + /* no-op */ + } +#endif + } + + static dim3 get_block_shape() { + dim3 block(MaxThreadsPerBlock, 1, 1); + return block; + } + + static dim3 get_grid_shape(Params const& params) { + auto [Q, K, D, D_VO, HB] = params.problem_shape; + auto [H, B] = HB; + auto [H_R, H_K] = H; + dim3 grid(ceil_div(K, TileShapeK{}), H_K, B); + return grid; + } +}; + +} // namespace cutlass::fmha::kernel diff --git a/natten/csrc/include/natten/cuda/fmha_blackwell/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp b/natten/csrc/include/natten/cuda/fmha_blackwell/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp new file mode 100644 index 00000000..d067ec81 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_blackwell/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp @@ -0,0 +1,757 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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/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 "natten/cuda/fmha_blackwell/collective/fmha_common.hpp" +#include "natten/cuda/fmha_blackwell/collective/fmha_fusion.hpp" +#include "natten/cuda/fmha_blackwell/kernel/fmha_causal_tile_scheduler.hpp" +#include "natten/cuda/fmha_blackwell/kernel/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; +}; + +struct Sm100MlaFwdCtxKernelWarpspecializedSchedule { + 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 = 184; + static const int NumRegsCorrection = 96 - (kDebugUsingPrintf ? 16 : 0); + static const int NumRegsOther = 48 + (kDebugUsingPrintf ? 16 : 0); + static const int NumRegsEmpty = 24; + + static const int NumWarps = 16; +}; + +template < + class ProblemShapeIn, + class CollectiveMainloop, + class CollectiveEpilogue, + class TileScheduler, + class KernelSchedule = Sm100FmhaCtxKernelWarpspecializedSchedule> +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_assert(NumWarpsEpilogue == CollectiveEpilogue::NumWarpsEpilogue); + static_assert(NumWarpsLoad == CollectiveEpilogue::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; + + static constexpr bool IsMla = std:: + is_same_v; + + using ClusterShape = typename CollectiveMainloop::ClusterShape; + + using TmemAllocator = cute::TMEM::Allocator1Sm; + + struct SharedStorage { + using UnionType = union { + typename CollectiveMainloop::TensorStorage mainloop; + typename CollectiveEpilogue::TensorStorage epilogue; + }; + + using StructType = struct { + typename CollectiveMainloop::TensorStorage mainloop; + typename CollectiveEpilogue::TensorStorage epilogue; + }; + + static constexpr bool IsPersistent = + std::is_same_v || + std::is_same_v; + using MainloopEpilogueStorage = std::conditional_t< + IsPersistent, + std::conditional_t< + IsMla, + std::conditional_t< + CollectiveMainloop::IsOrderLoadEpilogue, + UnionType, + StructType>, + StructType>, + UnionType>; + + MainloopEpilogueStorage mainloop_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); + + struct Arguments { + ProblemShape problem_shape; + typename CollectiveMainloop::Arguments mainloop; + typename CollectiveEpilogue::Arguments epilogue; + 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.problem_shape, args.hw_info, ClusterShape{}, TileShape{})}; + } + + 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 ( \ + !defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) && \ + !defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) && \ + !defined(CUTLASS_ARCH_MMA_SM103A_ENABLED) && \ + !defined(CUTLASS_ARCH_MMA_SM103F_ENABLED)) + printf( + "ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); +#else + + 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); + + auto get_epilogue_storage = [&]() { + if constexpr (IsMla && CollectiveMainloop::IsOrderLoadEpilogue) { + return reinterpret_cast( + shared_storage.mainloop_epilogue.mainloop.smem_o.data()); + } else { + return &shared_storage.mainloop_epilogue.epilogue; + } + }; + typename CollectiveEpilogue::TensorStorage& epilogue_storage = + *get_epilogue_storage(); + + 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::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.is_leader = + lane_predicate && (role == WarpRole::Load); + pipeline_load_kv_params.transaction_bytes = + CollectiveMainloop::TransactionBytesLoadK; + 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::PipelineQ>(); + + 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::PipelineKV>(); + + 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>(); + + 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::PipelineS>(); + + 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>(); + + 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::PipelineC>(); + + 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::PipelineE>(); + + typename CollectiveMainloop::PipelineO::PipelineState + pipeline_mma_corr_consumer_state; + typename CollectiveMainloop::PipelineO::PipelineState + pipeline_mma_corr_producer_state = cutlass::make_producer_start_state< + typename CollectiveMainloop::PipelineO>(); + + 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) { + 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(); + + bool has_valid = false; + + 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; + } + + has_valid = true; + + if (get<1>(logical_problem_shape) == 0) { + mainloop.correction_empty( + blk_coord, + params.mainloop, + logical_problem_shape, + params.problem_shape, + epilogue_storage, + pipeline_corr_epi, + pipeline_corr_epi_producer_state, + epilogue); + continue; + } + + mainloop.correction( + blk_coord, + params.mainloop, + logical_problem_shape, + params.problem_shape, + epilogue_storage, + 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); + + if (has_valid) { + 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(); + + bool allocated = false; + + 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 (!allocated) { + tmem_allocator.allocate( + TmemAllocator::Sm100TmemCapacityColumns, + &shared_storage.tmem_base_ptr); + __syncwarp(); + allocated = true; + } + + if (get<1>(logical_problem_shape) == 0) { + continue; + } + + mainloop.mma( + blk_coord, + params.mainloop, + logical_problem_shape, + shared_storage.mainloop_epilogue.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(); + + if constexpr (IsMla && CollectiveMainloop::IsOrderLoadEpilogue) { + cutlass::arch::NamedBarrier::arrive( + (NumWarpsLoad + NumWarpsEpilogue) * NumThreadsPerWarp, + cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); + } + + 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) { + continue; + } + + mainloop.load( + blk_coord, + logical_problem_shape, + params.mainloop, + params.problem_shape, + shared_storage.mainloop_epilogue.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(); + + bool has_valid = false; + + 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; + } + + has_valid = true; + + epilogue.store( + blk_coord, + logical_problem_shape, + params.epilogue, + params.problem_shape, + epilogue_storage, + pipeline_corr_epi, + pipeline_corr_epi_consumer_state); + } + + static_assert(NumWarpsEpilogue <= 1); + if constexpr (NumWarpsEpilogue == 1) { + if (has_valid) { + 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 */ + } +#endif + } +}; + +} // namespace cutlass::fmha::kernel diff --git a/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_bwd_tma_warpspecialized.hpp b/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_bwd_tma_warpspecialized.hpp new file mode 100644 index 00000000..090733c7 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_bwd_tma_warpspecialized.hpp @@ -0,0 +1,1235 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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/gemm/collective/collective_builder.hpp" + +#include "natten/cuda/fmha_hopper/collective/fmha_collective_load.hpp" +#include "natten/cuda/fmha_hopper/collective/fmha_collective_softmax.hpp" +#include "natten/cuda/fmha_hopper/collective/fmha_common.hpp" +#include "natten/cuda/fmha_hopper/collective/fmha_varlen.hpp" +#include "natten/cuda/fmha_hopper/kernel/fmha_options.hpp" + +namespace cutlass::fmha::collective { + +template < + typename Element_, + typename ElementAccumulator_, + typename TileShape_, // BlockQO, BlockKV, BlockHead + class Fusion, + class... Options> +struct FmhaBwdMainloopTmaWarpSpecializedSm90 { + using Element = Element_; + using ElementAccumulator = ElementAccumulator_; + using TileShape = TileShape_; + + static constexpr bool kIsPersistent = false; + + static const int NumLoadWarpGroups = 1; + static constexpr int NumMmaWarpGroups = 2; + static constexpr int StageCountQ = 2 /*K, V*/ * NumMmaWarpGroups; + static constexpr int StageCount = 2 /*Q, dO*/ * 2 /* actual stages */; + + static const int kOuterLoads = 2; + using StagesQ = cutlass::gemm::collective::StageCount; + using Stages = cutlass::gemm::collective::StageCount; + using ClusterShape = Shape<_1, _1, _1>; + static_assert(StagesQ::value >= 2); + static_assert(Stages::value >= 2 * NumMmaWarpGroups); + + // 16B alignment lets us use TMA + static constexpr int Alignment = 16 / sizeof(Element); + + using TileShapeNM = Shape< // (N,M,D) + decltype(tuple_element_t<1, TileShape>{} / Int{}), + tuple_element_t<0, TileShape>, + tuple_element_t<2, TileShape>>; + + using TileShapeND = decltype(select<0, 2, 1>(TileShapeNM{})); // (N,D,M) + + using TileShapeMD = decltype(select<2, 1, 0>(TileShapeND{})); // (M,D,N) + + using CollectiveMmaNM = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm90, + cutlass::arch::OpClassTensorOp, + Element, + cute::tuple>, + Alignment, + Element, + cute::tuple>, + Alignment, + ElementAccumulator, + TileShapeNM, + ClusterShape, + Stages, + cutlass::gemm::KernelTmaWarpSpecialized>::CollectiveOp; + + using CollectiveMmaND = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm90, + cutlass::arch::OpClassTensorOp, + Element, + cute::tuple>, + Alignment, // from register, doesn't matter + Element, + cute::tuple<_1, int, cute::tuple>, + Alignment, + ElementAccumulator, + TileShapeND, + ClusterShape, + Stages, + cutlass::gemm::KernelTmaWarpSpecialized>::CollectiveOp; + + using CollectiveMmaND_SS = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm90, + cutlass::arch::OpClassTensorOp, + Element, + cute::tuple>, + Alignment, // from register, doesn't matter + Element, + cute::tuple<_1, int, cute::tuple>, + Alignment, + ElementAccumulator, + TileShapeND, + ClusterShape, + Stages, + cutlass::gemm::KernelTmaWarpSpecialized>::CollectiveOp; + + using CollectiveMmaMD = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm90, + cutlass::arch::OpClassTensorOp, + Element, + cute::tuple<_1, int, cute::tuple>, + Alignment, // from smem, might matter (?) + Element, + cute::tuple<_1, int, cute::tuple>, + Alignment, + ElementAccumulator, + TileShapeMD, + ClusterShape, + Stages, + cutlass::gemm::KernelTmaWarpSpecialized>::CollectiveOp; + + using TiledMmaNM = typename CollectiveMmaNM::TiledMma; + using TiledMmaND_SS = typename CollectiveMmaND_SS::TiledMma; + using TiledMmaND_RS = + decltype(convert_to_gmma_rs(typename CollectiveMmaND::TiledMma{})); + using TiledMmaND = TiledMmaND_RS; + using TiledMmaMD = typename CollectiveMmaMD::TiledMma; + + using SmemLayoutQ = typename CollectiveMmaNM::SmemLayoutB; + using SmemLayoutK = typename CollectiveMmaNM::SmemLayoutA; + using SmemLayoutV = typename CollectiveMmaNM::SmemLayoutA; + using SmemLayoutDO = typename CollectiveMmaNM::SmemLayoutB; + + // using SmemLayoutDQ = Layout< + // Shape< + // tuple_element_t<0, TileShapeMD>, + // Shape<_2, _4, decltype(tuple_element_t<1, TileShapeMD>{} / _8{})>, + // _2 + // >, + // Stride< + // _4, + // Stride{} * _4{}), _1, + // decltype(tuple_element_t<0, TileShapeMD>{} * _8{})>, + // decltype(tuple_element_t<0, TileShapeMD>{} * tuple_element_t<1, + // TileShapeMD>{}) + // >>; + + using SmemLayoutDQ_0 = Layout< + Shape< + tuple_element_t<0, TileShapeMD>, + tuple_element_t<1, TileShapeMD>, + _2>, + Stride< + tuple_element_t<1, TileShapeMD>, + _1, + decltype(tuple_element_t<0, TileShapeMD>{} * tuple_element_t<1, TileShapeMD>{})>>; + + using SmemAtomDQ = + decltype(cutlass::gemm::collective::detail::ss_smem_selector< + cute::GMMA::Major::K, + ElementAccumulator, + tuple_element_t<0, TileShapeMD>, + tuple_element_t<1, TileShapeMD>>()); + using SmemLayoutDQ_1 = decltype(tile_to_shape( + SmemAtomDQ{}, + make_shape(get<0>(TileShapeMD{}), get<1>(TileShapeMD{}), _2{}), + Step<_2, _1, _3>{})); + using SmemLayoutDQ = SmemLayoutDQ_1; + + using PipelineDQ = cutlass::PipelineAsync<2>; + + using SmemLayoutDS_0 = decltype(unstageSmemLayout( + typename CollectiveMmaMD::SmemLayoutA{}, + Int{})); + + using SmemLayoutDS = decltype(tile_to_shape( + GMMA::Layout_MN_INTER_Atom{}, + make_shape( + size<0>(SmemLayoutDS_0{}), + size<1>(SmemLayoutDS_0{}), + size<2>(SmemLayoutDS_0{})), + Step<_1, _2, _3>{})); + using SmemLayoutKp = typename CollectiveMmaMD::SmemLayoutB; + + using SmemLayoutQp = typename CollectiveMmaND::SmemLayoutB; + using SmemLayoutDOp = typename CollectiveMmaND::SmemLayoutB; + + using SmemLayoutLSE = + Layout, Int>>; + + using MainloopPipeline = cutlass::PipelineTmaAsync; + using MainloopPipelineQ = cutlass::PipelineTmaAsync; + + using PipelineState = + typename cutlass::PipelineState; + using PipelineStateQ = + typename cutlass::PipelineState; + + using TileShapePV = TileShapeND; // To work with the kernel level + using TiledMmaPV = TiledMmaND; + + static constexpr int kInnerLoadBytes = + size(SmemLayoutQ{}(_, _, _0{})) * sizeof(Element) + + size(SmemLayoutLSE{}(_, _0{})) * sizeof(ElementAccumulator); + static constexpr int kOuterLoadBytes = + size(SmemLayoutK{}(_, _, _0{})) * sizeof(Element); + + struct SharedStorage { + // One for each consumer WG + union { + cute::array_aligned> smem_k; + cute::array_aligned> smem_kp; + cute::array_aligned> smem_v; + }; + + cute::array_aligned> smem_ds; + + // Loaded by producer, consumed by both WGs + union { + cute::array_aligned> smem_q; + cute::array_aligned> smem_do; + cute::array_aligned> smem_qp; + cute::array_aligned> smem_dop; + }; + + // Accumulated into by both consumers, potentially loaded, potentially + // written + cute::array_aligned> + smem_dq; + + union { + cute::array_aligned> + smem_lse; + cute::array_aligned> + smem_sumOdO; + }; + }; + + struct Arguments { + const Element* ptr_Q; + cute::tuple dQ; + const Element* ptr_K; + cute::tuple dK; + const Element* ptr_V; + cute::tuple dV; + + const Element* ptr_dO; + cute::tuple dDO; + + const ElementAccumulator* ptr_LSE; + cute::tuple dLSE; + const ElementAccumulator* ptr_sum_OdO; + cute::tuple dSumOdO; + + ElementAccumulator* ptr_dQ; + cute::tuple dDQ; + + // if zero, defaults to 1/sqrt(D) + float scale_softmax = 0.0f; + }; + + using TMA_Q = typename CollectiveMmaNM::Params::TMA_B; + using TMA_K = typename CollectiveMmaNM::Params::TMA_A; + using TMA_V = typename CollectiveMmaNM::Params::TMA_A; + using TMA_DO = typename CollectiveMmaNM::Params::TMA_B; + + using TMA_LSE = decltype(make_tma_copy( + SM90_TMA_LOAD{}, + make_tensor( + (const ElementAccumulator*)nullptr, + make_shape(1, 1, 1), + make_stride(_1{}, 0, 0)), + SmemLayoutLSE{}(_, _0{}))); + using TMA_ODO = TMA_LSE; + + using TMA_DQ = decltype(make_tma_copy( + SM90_TMA_REDUCE_ADD{}, + make_tensor( + (const ElementAccumulator*)nullptr, + make_shape(1, 1, 1, 1), + make_stride(0, _1{}, 0, 0)), + SmemLayoutDQ{}(_, _, _0{}))); + + using LoadQ = CollectiveLoadTma< + LoadKind::kBwdM, + MainloopPipeline, + Element, + SmemLayoutQ, + TMA_Q>; + + using LoadK = CollectiveLoadTma< + LoadKind::kBwdN, + MainloopPipelineQ, + Element, + SmemLayoutK, + TMA_K>; + + using LoadV = CollectiveLoadTma< + LoadKind::kBwdN, + MainloopPipelineQ, + Element, + SmemLayoutV, + TMA_V>; + + using LoadDO = CollectiveLoadTma< + LoadKind::kBwdM, + MainloopPipeline, + Element, + SmemLayoutDO, + TMA_DO>; + + using LoadLSE = CollectiveLoadTma< + LoadKind::kBwdScalar, + MainloopPipeline, + ElementAccumulator, + SmemLayoutLSE, + TMA_LSE>; + + using LoadODO = CollectiveLoadTma< + LoadKind::kBwdScalar, + MainloopPipeline, + ElementAccumulator, + SmemLayoutLSE, + TMA_ODO>; + + struct Params { + TMA_Q tma_load_q; + TMA_K tma_load_k; + TMA_V tma_load_v; + TMA_DO tma_load_do; + + TMA_LSE tma_load_lse; + TMA_ODO tma_load_odo; + + TMA_DQ tma_red_dq; + + float scale_softmax; + float scale_softmax_log2; + }; + + static_assert(size(TiledMmaNM{}) == size(TiledMmaND{})); + static_assert(size(TiledMmaNM{}) == size(TiledMmaMD{})); + + template + static bool can_implement( + ProblemShape const& problem_size, + Arguments const& args) { + // NOTE (ali): the original Hopper FMHA kernel does an alignment check on + // seqlen_Q as well, citing TMA requirement. However, it does not look like + // TMA runs into any issues, because cases that violate this run perfectly + // fine. + // + // Backprop requires this extra check though, and the point of failure is + // LSE. + // + // This does not cause any issues with FNA kernels, since they already pad + // for tile size, and tile sizes are usually powers of two, and at least 64, + // so divisibility by the required alignment 16 is a guarantee, though not + // explicitly asserted. + // + // Adding in padding along seqlen_Q requires modifications in the bwd mask + // as well. + // + // Therefore, we only pad LSE, and not any other tensors to get around this + // issue in bwd, instead of padding Q, dQ, O and dO as well. + // + // If we pad those, we would have to modify the residual mask as well. + // + return true && (get<4>(problem_size) <= get<2>(TileShape{})) && + ((get<4>(problem_size) % Alignment) == 0) + /*&& ((get<2>(problem_size) % Alignment) == 0)*/; + } + + template + static Params to_underlying_arguments( + ProblemShape const& problem_size, + Arguments const& args, + void* workspace) { + auto problem_shape_nm = apply_variable_length_scheduler(make_shape( + get<3>(problem_size), + get<2>(problem_size), + get<4>(problem_size), + make_shape(get<0>(problem_size), get<1>(problem_size)))); + + auto dK = make_stride( + get<2>(args.dK), + get<3>(args.dK), + make_stride(get<0>(args.dK), get<1>(args.dK))); + auto dQ = make_stride( + get<2>(args.dQ), + get<3>(args.dQ), + make_stride(get<0>(args.dQ), get<1>(args.dQ))); + auto params_nm_kq = CollectiveMmaNM::to_underlying_arguments( + problem_shape_nm, + typename CollectiveMmaNM::Arguments{ + args.ptr_K, + dK, + args.ptr_Q, + dQ, + }, + /*workspace=*/nullptr); + + auto dV = make_stride( + get<2>(args.dV), + get<3>(args.dV), + make_stride(get<0>(args.dV), get<1>(args.dV))); + auto dDO = make_stride( + get<2>(args.dDO), + get<3>(args.dDO), + make_stride(get<0>(args.dDO), get<1>(args.dDO))); + auto params_nm_vdo = CollectiveMmaNM::to_underlying_arguments( + problem_shape_nm, + typename CollectiveMmaNM::Arguments{ + args.ptr_V, + dV, + args.ptr_dO, + dDO, + }, + /*workspace=*/nullptr); + + // These are batch-packed, not sequence-packed, so they should use max + // seqlen (default) over total. + TMA_LSE tma_load_lse = make_tma_copy( + SM90_TMA_LOAD{}, + make_tensor( + args.ptr_LSE, + select<2, 0, 1>(problem_size), + select<2, 0, 1>(args.dLSE)), + SmemLayoutLSE{}(_, _0{})); + TMA_ODO tma_load_odo = make_tma_copy( + SM90_TMA_LOAD{}, + make_tensor( + args.ptr_sum_OdO, + select<2, 0, 1>(problem_size), + select<2, 0, 1>(args.dSumOdO)), + SmemLayoutLSE{}(_, _0{})); + + TMA_DQ tma_red_dq = make_tma_copy( + SM90_TMA_REDUCE_ADD{}, + make_tensor( + args.ptr_dQ, + select<2, 4, 0, 1>(problem_size), + select<2, 3, 0, 1>(args.dDQ)), + SmemLayoutDQ{}(_, _, _0{})); + + float scale_softmax = args.scale_softmax; + if (scale_softmax == 0.0f) { + scale_softmax = 1.0f / (float)std::sqrt(get<4>(problem_size)); + } + float log2_e = static_cast(std::log2(std::exp(1.0))); + + return Params{ + params_nm_kq.tma_load_b, + params_nm_kq.tma_load_a, + params_nm_vdo.tma_load_a, + params_nm_vdo.tma_load_b, + tma_load_lse, + tma_load_odo, + tma_red_dq, + scale_softmax, + log2_e * scale_softmax, + }; + } + + 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()); + cute::prefetch_tma_descriptor(params.tma_load_do.get_tma_descriptor()); + cute::prefetch_tma_descriptor(params.tma_load_odo.get_tma_descriptor()); + cute::prefetch_tma_descriptor(params.tma_load_lse.get_tma_descriptor()); + } + + template < + bool kLoadOuter, + class BlkCoord, + class ProblemShape, + class LoadWarpBarrier> + CUTLASS_DEVICE void load_kv_maybe_q( + int block_rank_in_cluster, + BlkCoord const& blk_coord, + Params const& params, + ProblemShape const& problem_size, + MainloopPipeline& pipeline_inner, + PipelineState& smem_pipe_write_inner, + MainloopPipelineQ& pipeline_outer, + PipelineStateQ& smem_pipe_write_outer, + SharedStorage& storage, + LoadWarpBarrier& load_warp_barrier, + bool do_barrier) { + // Load pattern: + // K0 V0 K1 V1 + // Q0 DO0 Q1 DO1 Q2 DO2 ... + // K0 Q0 V0 K1 DO0 V1 ... + int lane_predicate = cute::elect_one_sync(); + + int outer_tile_count = NumMmaWarpGroups; + auto problem_size_f = + apply_variable_length(problem_size, get<2, 0>(blk_coord)); + int inner_tile_count = + Fusion{}.get_trip_count(blk_coord, TileShape{}, problem_size_f); + + auto outer_tile_iter = cute::make_coord_iterator(outer_tile_count); + auto inner_tile_iter = cute::make_coord_iterator(inner_tile_count); + + uint16_t mcast_mask_b = 0; + + LoadQ load_q{params.tma_load_q, pipeline_inner, storage.smem_q}; + auto load_state_q = load_q.init_state( + block_rank_in_cluster, + problem_size, + TileShapeNM{}, + blk_coord, + inner_tile_count); + + LoadDO load_do{params.tma_load_do, pipeline_inner, storage.smem_do}; + auto load_state_do = load_do.init_state( + block_rank_in_cluster, + problem_size, + TileShapeNM{}, + blk_coord, + inner_tile_count); + + LoadK load_k{params.tma_load_k, pipeline_outer, storage.smem_k}; + auto load_state_k = load_k.init_state( + _0{}, problem_size, TileShapeNM{}, blk_coord, outer_tile_count); + + LoadV load_v{params.tma_load_v, pipeline_outer, storage.smem_v}; + auto load_state_v = load_v.init_state( + _0{}, problem_size, TileShapeNM{}, blk_coord, outer_tile_count); + + LoadLSE load_lse{params.tma_load_lse, pipeline_inner, storage.smem_lse}; + auto load_state_lse = load_lse.init_state( + _0{}, problem_size, TileShapeNM{}, blk_coord, outer_tile_count); + + LoadODO load_odo{params.tma_load_odo, pipeline_inner, storage.smem_sumOdO}; + auto load_state_odo = load_odo.init_state( + _0{}, problem_size, TileShapeNM{}, blk_coord, outer_tile_count); + + outer_tile_count *= 2; // K & V + inner_tile_count *= 4; // Q & dO & LSE & sumOdO + + while (inner_tile_count > 0) { + if (Fusion{}.is_contributing( + make_coord(*inner_tile_iter, get<1>(blk_coord)), + TileShape{}, + problem_size_f)) { + break; + } + inner_tile_count -= 4; + ++inner_tile_iter; + } + + if constexpr (kLoadOuter) { + load_k.template step( + outer_tile_iter, + load_state_k, + smem_pipe_write_outer, + lane_predicate, + outer_tile_count); + } + + load_q.template step( + inner_tile_iter, + load_state_q, + smem_pipe_write_inner, + lane_predicate, + inner_tile_count, + mcast_mask_b); + load_lse.template step( + inner_tile_iter, + load_state_lse, + smem_pipe_write_inner, + lane_predicate, + inner_tile_count, + mcast_mask_b); + + if constexpr (!kLoadOuter) { + if (do_barrier) { + load_warp_barrier.arrive(); + load_warp_barrier.wait(/*phase=*/0); + do_barrier = false; + } + } + + if constexpr (kLoadOuter) { + load_v.template step( + outer_tile_iter, + load_state_v, + smem_pipe_write_outer, + lane_predicate, + outer_tile_count); + load_k.template step( + outer_tile_iter, + load_state_k, + smem_pipe_write_outer, + lane_predicate, + outer_tile_count); + } + + load_do.template step( + inner_tile_iter, + load_state_do, + smem_pipe_write_inner, + lane_predicate, + inner_tile_count, + mcast_mask_b); + load_odo.template step( + inner_tile_iter, + load_state_odo, + smem_pipe_write_inner, + lane_predicate, + inner_tile_count, + mcast_mask_b); + + if constexpr (kLoadOuter) { + load_v.template step( + outer_tile_iter, + load_state_v, + smem_pipe_write_outer, + lane_predicate, + outer_tile_count); + } + + if constexpr (kLoadOuter) { + while (outer_tile_count > 0) { + load_k.template step( + outer_tile_iter, + load_state_k, + smem_pipe_write_outer, + lane_predicate, + outer_tile_count); + load_v.template step( + outer_tile_iter, + load_state_v, + smem_pipe_write_outer, + lane_predicate, + outer_tile_count); + } + } + + CUTLASS_PRAGMA_NO_UNROLL + while (inner_tile_count > 0) { + while (inner_tile_count > 0) { + if (Fusion{}.is_contributing( + make_coord(*inner_tile_iter, get<1>(blk_coord)), + TileShape{}, + problem_size_f)) { + break; + } + inner_tile_count -= 4; + ++inner_tile_iter; + } + load_q.template step( + inner_tile_iter, + load_state_q, + smem_pipe_write_inner, + lane_predicate, + inner_tile_count, + mcast_mask_b); + load_lse.template step( + inner_tile_iter, + load_state_lse, + smem_pipe_write_inner, + lane_predicate, + inner_tile_count, + mcast_mask_b); + + load_do.template step( + inner_tile_iter, + load_state_do, + smem_pipe_write_inner, + lane_predicate, + inner_tile_count, + mcast_mask_b); + load_odo.template step( + inner_tile_iter, + load_state_odo, + smem_pipe_write_inner, + lane_predicate, + inner_tile_count, + mcast_mask_b); + } + } + + template + CUTLASS_DEVICE void load_maybe_q( + BlkCoord const& blk_coord, + Params const& params, + ProblemShape const& problem_size, + MainloopPipelineQ& pipeline_outer, + PipelineStateQ& smem_pipe_write_outer, + SharedStorage& storage, + LoadWarpBarrier& load_warp_barrier, + bool do_barrier) { + // Load pattern: + // K0 V0 K1 V1 + // Q0 DO0 Q1 DO1 Q2 DO2 ... + // K0 Q0 V0 K1 DO0 V1 ... + int lane_predicate = cute::elect_one_sync(); + + int outer_tile_count = NumMmaWarpGroups; + + auto outer_tile_iter = cute::make_coord_iterator(outer_tile_count); + + LoadK load_k{params.tma_load_k, pipeline_outer, storage.smem_k}; + auto load_state_k = load_k.init_state( + _0{}, problem_size, TileShapeNM{}, blk_coord, outer_tile_count); + + LoadV load_v{params.tma_load_v, pipeline_outer, storage.smem_v}; + auto load_state_v = load_v.init_state( + _0{}, problem_size, TileShapeNM{}, blk_coord, outer_tile_count); + + outer_tile_count *= 2; // K & V + + load_k.template step( + outer_tile_iter, + load_state_k, + smem_pipe_write_outer, + lane_predicate, + outer_tile_count); + + if (do_barrier) { + load_warp_barrier.arrive(); + load_warp_barrier.wait(/*phase=*/0); + do_barrier = false; + } + + load_v.template step( + outer_tile_iter, + load_state_v, + smem_pipe_write_outer, + lane_predicate, + outer_tile_count); + + while (outer_tile_count > 0) { + load_k.template step( + outer_tile_iter, + load_state_k, + smem_pipe_write_outer, + lane_predicate, + outer_tile_count); + load_v.template step( + outer_tile_iter, + load_state_v, + smem_pipe_write_outer, + lane_predicate, + outer_tile_count); + } + } + + template < + class BlkCoord, + class ProblemShape, + class MainloopPipelineReducer, + class PipelineStateReducer> + CUTLASS_DEVICE void reduce( + BlkCoord const& blk_coord, + Params const& params, + ProblemShape const& problem_size, + MainloopPipelineReducer& pipeline_reducer, + PipelineStateReducer& smem_pipe_read_reducer, + SharedStorage& storage) { + int lane_predicate = cute::elect_one_sync(); + + Tensor mDQ_full = + params.tma_red_dq.get_tma_tensor(select<2, 4, 0, 1>(problem_size)); + Tensor gDQ_full = local_tile( + mDQ_full, + TileShapeMD{}, + make_coord(_, _, _), + Step<_1, _1, Underscore>{}); + Tensor gDQ = + gDQ_full(_, _, _, _0{}, get<2, 0>(blk_coord), get<2, 1>(blk_coord)); + Tensor sDQ = + make_tensor(make_smem_ptr(storage.smem_dq.data()), SmemLayoutDQ{}); + + auto block_tma = params.tma_red_dq.get_slice(_0{}); + + Tensor tDQsDQ = block_tma.partition_S(sDQ); + Tensor tDQgDQ = block_tma.partition_D(gDQ); + + auto problem_size_f = + apply_variable_length(problem_size, get<2, 0>(blk_coord)); + int inner_tile_count = + Fusion{}.get_trip_count(blk_coord, TileShape{}, problem_size_f); + int g_index = 0; + + auto smem_pipe_release_reducer = smem_pipe_read_reducer; + bool first = true; + while (inner_tile_count > 0) { + while (inner_tile_count > 0) { + if (Fusion{}.is_contributing( + make_coord(g_index, get<1>(blk_coord)), + TileShape{}, + problem_size_f)) { + break; + } + inner_tile_count -= 1; + ++g_index; + } + if (inner_tile_count == 0) + break; + + pipeline_reducer.consumer_wait(smem_pipe_read_reducer); + if (lane_predicate == 1) { + tma_store_wait<1>(); + } + if (!first) { + pipeline_reducer.consumer_release(smem_pipe_release_reducer); + ++smem_pipe_release_reducer; + } else { + first = false; + } + if (lane_predicate == 1) { + copy( + params.tma_red_dq, + tDQsDQ(_, _, _, smem_pipe_read_reducer.index()), + tDQgDQ(_, _, _, g_index)); + tma_store_arrive(); + } + ++smem_pipe_read_reducer; + --inner_tile_count; + ++g_index; + } + if (lane_predicate) { + tma_store_wait<0>(); + } + pipeline_reducer.consumer_release(smem_pipe_release_reducer); + ++smem_pipe_release_reducer; + } + + template < + class BlkCoord, + class ProblemShape, + class MainloopPipelineReducer, + class PipelineStateReducer, + class MathWgOrderBarrier> + CUTLASS_DEVICE auto compute( + BlkCoord const& blk_coord, + BlkCoord const& wg_coord, + Params const& params, + ProblemShape const& problem_size, + MainloopPipeline& pipeline_inner, + PipelineState& smem_pipe_read_inner, + MainloopPipelineQ& pipeline_outer, + PipelineStateQ& smem_pipe_read_outer, + MainloopPipelineReducer& pipeline_reducer, + PipelineStateReducer& smem_pipe_write_reducer, + SharedStorage& storage, + MathWgOrderBarrier& math_wg_order_barrier) { + TiledMmaND tiled_mma_nd; + + Tensor acc_DV = + partition_fragment_C(tiled_mma_nd, take<0, 2>(TileShapeND{})); + clear(acc_DV); + + Tensor acc_DK = + partition_fragment_C(tiled_mma_nd, take<0, 2>(TileShapeND{})); + clear(acc_DK); + + int thread_idx = int(threadIdx.x) % cutlass::NumThreadsPerWarpGroup; + + PipelineState smem_pipe_release_inner = smem_pipe_read_inner; + + pipeline_outer.consumer_wait(smem_pipe_read_outer); + PipelineStateQ smem_pipe_read_k = smem_pipe_read_outer; + ++smem_pipe_read_outer; + pipeline_outer.consumer_wait(smem_pipe_read_outer); + PipelineStateQ smem_pipe_read_v = smem_pipe_read_outer; + + auto problem_size_f = + apply_variable_length(problem_size, get<2, 0>(blk_coord)); + int inner_tile_count = + Fusion{}.get_trip_count(blk_coord, TileShape{}, problem_size_f); + + TiledMmaNM tiled_mma_nm; + Tensor sK = + make_tensor(make_smem_ptr(storage.smem_k.data()), SmemLayoutK{}); + Tensor sQ = + make_tensor(make_smem_ptr(storage.smem_q.data()), SmemLayoutQ{}); + auto thr_mma_nm = tiled_mma_nm.get_thread_slice(thread_idx); + Tensor tSsK = thr_mma_nm.partition_A(sK); + Tensor tSsQ = thr_mma_nm.partition_B(sQ); + Tensor tSrK = thr_mma_nm.make_fragment_A(tSsK); + Tensor tSrQ = thr_mma_nm.make_fragment_B(tSsQ); + + Tensor sV = + make_tensor(make_smem_ptr(storage.smem_v.data()), SmemLayoutV{}); + Tensor sDO = + make_tensor(make_smem_ptr(storage.smem_do.data()), SmemLayoutDO{}); + + Tensor tDPsV = thr_mma_nm.partition_A(sV); + Tensor tDPsDO = thr_mma_nm.partition_B(sDO); + Tensor tDPrV = thr_mma_nm.make_fragment_A(tDPsV); + Tensor tDPrDO = thr_mma_nm.make_fragment_B(tDPsDO); + + auto thr_mma_nd = tiled_mma_nd.get_thread_slice(thread_idx); + + Tensor sDOp = + make_tensor(make_smem_ptr(storage.smem_dop.data()), SmemLayoutDOp{}); + Tensor tDV_sDO = thr_mma_nd.partition_B(sDOp); + Tensor tDVrDO = thr_mma_nd.make_fragment_B(tDV_sDO); + + Tensor sQp = + make_tensor(make_smem_ptr(storage.smem_qp.data()), SmemLayoutQp{}); + Tensor tDK_sQ = thr_mma_nd.partition_B(sQp); + Tensor tDKrQ = thr_mma_nd.make_fragment_B(tDK_sQ); + + int wg_idx = + __shfl_sync(0xffffffff, get<1>(wg_coord) % NumMmaWarpGroups, 0); + + TiledMmaMD tiled_mma_md; + auto thr_mma_md = tiled_mma_md.get_thread_slice(thread_idx); + Tensor sDS = + make_tensor(make_smem_ptr(storage.smem_ds.data()), SmemLayoutDS{}); + Tensor tDQsDS = thr_mma_md.partition_A(sDS); + Tensor tDQrDS_full = thr_mma_md.make_fragment_A(tDQsDS); + Tensor tDQrDS = tDQrDS_full(_, _, _, _); + Tensor sKp = + make_tensor(make_smem_ptr(storage.smem_kp.data()), SmemLayoutKp{}); + Tensor tDQsK = thr_mma_md.partition_B(sKp); + Tensor tDQrK = thr_mma_md.make_fragment_B(tDQsK); + + Tensor sLSE = make_tensor( + make_smem_ptr(storage.smem_lse.data()), + make_shape( + get<0>(TileShapeNM{}), get<1>(TileShapeNM{}), Int{}), + make_stride(_0{}, _1{}, get<1>(TileShapeNM{}))); + Tensor tSsLSE = thr_mma_nm.partition_C(sLSE); + + Tensor sODO = make_tensor( + make_smem_ptr(storage.smem_sumOdO.data()), + make_shape( + get<0>(TileShapeNM{}), get<1>(TileShapeNM{}), Int{}), + make_stride(_0{}, _1{}, get<1>(TileShapeNM{}))); + Tensor tDPsODO = thr_mma_nm.partition_C(sODO); + + Tensor cS = make_identity_tensor(take<0, 2>(TileShapeNM{})); + Tensor tScS = thr_mma_nm.partition_C(cS); + int n_block = get<1>(wg_coord); + tScS.data() = tScS.data() + E<0>{} * n_block * get<0>(TileShapeNM{}); + + // Transpose + Tensor sDSp_full = sDS.compose(make_layout( + make_shape(size<1>(sDS), size<0>(sDS), size<2>(sDS)), + make_stride(size<0>(sDS), _1{}, size<1>(sDS) * size<0>(sDS)))); + Tensor sDSp = sDSp_full(_, _, _); + Tensor tDPsDS = thr_mma_nm.partition_C(sDSp); + + auto thr_mma_nd_ss = TiledMmaND_SS{}.get_thread_slice(thread_idx); + Tensor tDKsDSp = thr_mma_nd_ss.partition_A(sDSp); + + Tensor tDKrDSp = thr_mma_nd_ss.make_fragment_A(tDKsDSp); + + Tensor sDQ = + make_tensor(make_smem_ptr(storage.smem_dq.data()), SmemLayoutDQ{}); + auto tDQsDQ_full = thr_mma_md.partition_C(sDQ); + + auto smem_pipe_read_k_other = smem_pipe_read_k; + smem_pipe_read_k_other.advance(2); + + int k_index = 0; + + bool no_inner_tiles = true; + while (inner_tile_count > 0) { + while (inner_tile_count > 0) { + if (Fusion{}.is_contributing( + make_coord(k_index, get<1>(blk_coord)), + TileShape{}, + problem_size_f)) { + break; + } + inner_tile_count -= 1; + tScS.data() = tScS.data() + E<1>{} * get<1>(TileShapeNM{}); + k_index += 1; + } + if (inner_tile_count == 0) + break; + no_inner_tiles = false; + + pipeline_inner.consumer_wait(smem_pipe_read_inner); + PipelineState smem_pipe_read_q = smem_pipe_read_inner; + ++smem_pipe_read_inner; + PipelineState smem_pipe_read_do = smem_pipe_read_inner; + ++smem_pipe_read_inner; + + // GEMM KQ -> S + Tensor acc_S = + partition_fragment_C(tiled_mma_nm, take<0, 2>(TileShapeNM{})); + + warpgroup_fence_operand(acc_S); + warpgroup_arrive(); + gemm_zero_acc( + tiled_mma_nm, + tSrK(_, _, _, smem_pipe_read_k.index()), + tSrQ(_, _, _, smem_pipe_read_q.index()), + acc_S); + warpgroup_commit_batch(); + + pipeline_inner.consumer_wait(smem_pipe_read_do); + + // GEMM VdO -> dP + Tensor acc_DP = + partition_fragment_C(tiled_mma_nm, take<0, 2>(TileShapeNM{})); + + warpgroup_fence_operand(acc_DP); + warpgroup_arrive(); + gemm_zero_acc( + tiled_mma_nm, + tDPrV(_, _, _, smem_pipe_read_v.index()), + tDPrDO(_, _, _, smem_pipe_read_do.index()), + acc_DP); + warpgroup_commit_batch(); + + Tensor reg_LSE = make_fragment_like(acc_S); + for (int i = 0; i < size(reg_LSE); i++) { + reg_LSE(i) = ((ElementAccumulator)std::log2(std::exp(1.0))) * + tSsLSE(_, _, _, smem_pipe_read_q.index())(i); + } + + Tensor reg_ODO = make_fragment_like(acc_S); + if constexpr (decltype(get<0>(TileShape{}) != _128{})::value) { + for (int i = 0; i < size(reg_ODO); i++) { + reg_ODO(i) = tDPsODO(_, _, _, smem_pipe_read_do.index())(i); + } + } + + warpgroup_wait<1>(); + warpgroup_fence_operand(acc_S); + + math_wg_order_barrier.wait(); + // Compute S -> P + Fusion{}.before_softmax( + acc_S, tScS, problem_size_f); + auto acc_P = make_fragment_like(acc_S); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(acc_P); i++) { + acc_P(i) = ::exp2f(params.scale_softmax_log2 * acc_S(i) - reg_LSE(i)); + } + math_wg_order_barrier.arrive(); + + if constexpr (decltype(get<0>(TileShape{}) == _128{})::value) { + for (int i = 0; i < size(reg_ODO); i++) { + reg_ODO(i) = tDPsODO(_, _, _, smem_pipe_read_do.index())(i); + } + } + + warpgroup_wait<0>(); + warpgroup_fence_operand(acc_DP); + + // Compute dP P -> dS + auto acc_DS = make_fragment_like(acc_DP); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(acc_DS); i++) { + // We could move the scale out and into the respective epilogues (or a + // final scaling step) + acc_DS(i) = static_cast( + acc_P(i) * params.scale_softmax * (acc_DP(i) - reg_ODO(i))); + } + + // GEMM PdO -> dV + auto op_P = + make_acc_into_op(acc_P, typename TiledMmaND::LayoutA_TV{}); + warpgroup_fence_operand(acc_DV); + warpgroup_fence_operand(op_P); + warpgroup_arrive(); + cute::gemm( + tiled_mma_nd, + op_P, + tDVrDO(_, _, _, smem_pipe_read_do.index()), + acc_DV); + warpgroup_commit_batch(); + + // Store dS to smem dS' + if (wg_idx == 0) + math_wg_order_barrier.wait(); + + auto recast_bits = [](auto sz, auto t) { + return recast>(t); + }; + auto tDPsDS_v = recast_bits(Int * 2>{}, tDPsDS); + auto acc_DS_v = recast_bits(Int * 2>{}, acc_DS); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(acc_DS_v); i++) { + tDPsDS_v(_, _, _, wg_idx)(i) = acc_DS_v(i); + } + + cutlass::arch::fence_view_async_shared(); + if (wg_idx == 0) + math_wg_order_barrier.arrive(); + + // GEMM dS Q -> dK + if (wg_idx == 1) { + math_wg_order_barrier.wait(); + + // GEMM dS' K -> dQ + Tensor acc_DQ = + partition_fragment_C(tiled_mma_md, take<0, 2>(TileShapeMD{})); + + warpgroup_fence_operand(acc_DQ); + warpgroup_arrive(); + gemm_zero_acc( + tiled_mma_md, + tDQrDS(_, _, _, 0), + tDQrK(_, _, _, smem_pipe_read_k_other.index()), + acc_DQ); + cute::gemm( + tiled_mma_md, + tDQrDS(_, _, _, 1), + tDQrK(_, _, _, smem_pipe_read_k.index()), + acc_DQ); + warpgroup_commit_batch(); + + warpgroup_fence_operand(acc_DK); + warpgroup_arrive(); + cute::gemm( + TiledMmaND_SS{}, + tDKrDSp(_, _, _, wg_idx), + tDKrQ(_, _, _, smem_pipe_read_q.index()), + acc_DK); + warpgroup_commit_batch(); + + warpgroup_wait<1>(); + warpgroup_fence_operand(acc_DK); + + warpgroup_wait<1>(); + warpgroup_fence_operand(acc_DQ); + + math_wg_order_barrier.arrive(); + + pipeline_reducer.producer_acquire(smem_pipe_write_reducer); + auto tDQsDQ = tDQsDQ_full(_, _, _, smem_pipe_write_reducer.index()); + + // Store dQ to smem dQ' + // Invoke TMA reduce on dQ' + using Vec = uint_bit_t * 2>; + auto tDQsDQ_v = recast(tDQsDQ); + auto acc_DQ_v = recast(acc_DQ); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(acc_DQ_v); i++) { + tDQsDQ_v(i) = acc_DQ_v(i); + } + + cutlass::arch::fence_view_async_shared(); + + pipeline_reducer.producer_commit(smem_pipe_write_reducer); + ++smem_pipe_write_reducer; + } else { + warpgroup_fence_operand(acc_DK); + warpgroup_arrive(); + cute::gemm( + TiledMmaND_SS{}, + tDKrDSp(_, _, _, wg_idx), + tDKrQ(_, _, _, smem_pipe_read_q.index()), + acc_DK); + warpgroup_commit_batch(); + + warpgroup_wait<1>(); + warpgroup_fence_operand(acc_DK); + + pipeline_reducer.producer_acquire(smem_pipe_write_reducer); + pipeline_reducer.producer_commit(smem_pipe_write_reducer); + ++smem_pipe_write_reducer; + } + + --inner_tile_count; + + pipeline_inner.consumer_release(smem_pipe_release_inner); + ++smem_pipe_release_inner; + pipeline_inner.consumer_release(smem_pipe_release_inner); + ++smem_pipe_release_inner; + + tScS.data() = tScS.data() + E<1>{} * get<1>(TileShapeNM{}); + k_index += 1; + } + + if (no_inner_tiles) { + // when there's 0 inner tiles, we also need to make sure the order barrier + // stays in sync with the rest of the WG. + + // S -> P + math_wg_order_barrier.wait(); + math_wg_order_barrier.arrive(); + + math_wg_order_barrier.wait(); + math_wg_order_barrier.arrive(); + } + + pipeline_outer.consumer_release(smem_pipe_read_k); + pipeline_outer.consumer_release(smem_pipe_read_outer); + + // When there's 0 inner tiles (i.e. 1 Q tile, > 1 KV tiles, top-left causal + // mask has entire KV tiles to which nothing attends), reduce warp exits + // early and signals no arrivals on pipeline_reducer. If we don't jump over + // producer_tail, acquires will stall indefinitely. + if (not no_inner_tiles) { + pipeline_reducer.producer_tail(smem_pipe_write_reducer); + ++smem_pipe_read_outer; + } + + warpgroup_wait<0>(); + warpgroup_fence_operand(acc_DK); + warpgroup_fence_operand(acc_DV); + + return make_tuple(acc_DK, acc_DV); + } +}; + +} // namespace cutlass::fmha::collective diff --git a/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_load.hpp b/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_load.hpp new file mode 100644 index 00000000..c57da9a2 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_load.hpp @@ -0,0 +1,223 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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" + +#include "natten/cuda/fmha_hopper/collective/fmha_varlen.hpp" + +namespace cutlass::fmha::collective { + +enum class LoadKind { kQ, kK, kV, kBwdN, kBwdM, kBwdScalar }; + +template < + LoadKind kKind, + class Pipeline, + class Element, + class SmemLayout, + class TMA> +struct CollectiveLoadTma { + using Params = TMA; + using SharedStorage = + cute::array_aligned>; + using PipelineState = typename cutlass::PipelineState; + + Params const& params; + Pipeline& pipeline; + SharedStorage& storage; + + CUTLASS_DEVICE + CollectiveLoadTma( + Params const& params, + Pipeline& pipeline, + SharedStorage& storage) + : params(params), pipeline(pipeline), storage(storage) {} + + template + CUTLASS_DEVICE auto init_g( + ProblemSize const& problem_size, + TileShape const& tile_shape, + BlockCoord const& blk_coord_, + int loop_count) { + using X = Underscore; + int seq_off = 0; + BlockCoord blk_coord = blk_coord_; + if constexpr (is_problem_shape_variable_length(ProblemSize{})) { + int32_t* cumulative_length = nullptr; + + if constexpr ( + kKind == LoadKind::kK || kKind == LoadKind::kV || + kKind == LoadKind::kBwdN) { + cumulative_length = get<3>(problem_size).cumulative_length; + } else if constexpr (kKind == LoadKind::kQ || kKind == LoadKind::kBwdM) { + cumulative_length = get<2>(problem_size).cumulative_length; + } + + if (cumulative_length != nullptr) { + seq_off = cumulative_length[get<2, 0>(blk_coord)]; + get<2, 0>(blk_coord) = 0; + } + } + + if constexpr (kKind == LoadKind::kK) { + Tensor mK_tma = params.get_tma_tensor(make_shape( + get<3>(problem_size), + get<4>(problem_size), + select<0, 1>(problem_size))); + Tensor mK_full = domain_offset( + make_coord(seq_off, _0{}, make_coord(_0{}, _0{})), mK_tma); + Tensor gK_full = local_tile( + mK_full, tile_shape, make_coord(_, _, _), Step{}); + Tensor gK = gK_full(_, _, _, _0{}, get<2>(blk_coord)); + return gK; + } else if constexpr (kKind == LoadKind::kQ) { + Tensor mQ_tma = params.get_tma_tensor(make_shape( + get<2>(problem_size), + get<4>(problem_size), + select<0, 1>(problem_size))); + Tensor mQ_full = domain_offset( + make_coord(seq_off, _0{}, make_coord(_0{}, _0{})), mQ_tma); + Tensor gQ_full = local_tile( + mQ_full, tile_shape, make_coord(_, _, _), Step<_1, X, _1>{}); + Tensor gQ = gQ_full(_, _, _, _0{}, get<2>(blk_coord)); + return make_tensor( + gQ.data() + loop_count * get<0>(blk_coord) * stride<2>(gQ), + gQ.layout()); + } else if constexpr (kKind == LoadKind::kV) { + Tensor mV_tma = params.get_tma_tensor(make_shape( + get<4>(problem_size), + get<3>(problem_size), + select<0, 1>(problem_size))); + Tensor mV_full = domain_offset( + make_coord(_0{}, seq_off, make_coord(_0{}, _0{})), mV_tma); + Tensor gV_full = local_tile( + mV_full, tile_shape, make_coord(_, _, _), Step{}); + Tensor gV = gV_full(_, _, _0{}, _, get<2>(blk_coord)); + return gV; + } else if constexpr (kKind == LoadKind::kBwdN) { + Tensor m_tma = params.get_tma_tensor(make_shape( + get<3>(problem_size), + get<4>(problem_size), + select<0, 1>(problem_size))); + Tensor m_full = domain_offset( + make_coord(seq_off, _0{}, make_coord(_0{}, _0{})), m_tma); + Tensor g_full = local_tile( + m_full, tile_shape, make_coord(_, _, _), Step<_1, X, _1>{}); + Tensor g = g_full(_, _, _, _0{}, get<2>(blk_coord)); + return make_tensor( + g.data() + loop_count * get<1>(blk_coord) * stride<2>(g), g.layout()); + } else if constexpr (kKind == LoadKind::kBwdM) { + Tensor m_tma = params.get_tma_tensor(make_shape( + get<2>(problem_size), + get<4>(problem_size), + select<0, 1>(problem_size))); + Tensor m_full = domain_offset( + make_coord(seq_off, _0{}, make_coord(_0{}, _0{})), m_tma); + Tensor g_full = local_tile( + m_full, tile_shape, make_coord(_, _, _), Step{}); + Tensor g = g_full(_, _, _, _0{}, get<2>(blk_coord)); + return g; + } else if constexpr (kKind == LoadKind::kBwdScalar) { + // OdO and LSE are mapped to batch-packed intermediary tensors, which + // means they should offset along batch, not seq. + Tensor m_full = params.get_tma_tensor(select<2, 0, 1>(problem_size)); + Tensor g_full = + local_tile(m_full, tile_shape, make_coord(_, _, _), Step{}); + Tensor g = g_full(_, _, get<2, 0>(blk_coord_), get<2, 1>(blk_coord_)); + return g; + } + } + + template < + class ClusterRank, + class ProblemSize, + class TileShape, + class BlockCoord> + CUTLASS_DEVICE auto init_state( + ClusterRank const& block_rank_in_cluster, + ProblemSize const& problem_size, + TileShape const& tile_shape, + BlockCoord const& block_coord, + int loop_count) { + Tensor g = init_g(problem_size, tile_shape, block_coord, loop_count); + Tensor s = make_tensor(make_smem_ptr(storage.data()), SmemLayout{}); + + auto block_tma = params.get_slice(block_rank_in_cluster); + Tensor ts = block_tma.partition_D(s); + Tensor tg = block_tma.partition_S(g); + + return make_tuple(tg, ts); + } + + template < + bool kAdvanceIterator = true, + bool kAdvancePipe = true, + bool kAcquireBarrier = true, + class TileIterator, + class State> + CUTLASS_DEVICE void step( + TileIterator& tile_iter, + State const& state, + PipelineState& smem_pipe_write, + int lane_predicate, + int& tile_count, + uint16_t mcast_mask = 0) { + if ((lane_predicate == 1) && (tile_count > 0)) { + if constexpr (kAcquireBarrier) + pipeline.producer_acquire(smem_pipe_write); + using BarrierType = typename Pipeline::ProducerBarrierType; + BarrierType* tma_barrier = pipeline.producer_get_barrier(smem_pipe_write); + + if constexpr (kKind == LoadKind::kBwdScalar) { + copy( + params.with(*tma_barrier, mcast_mask), + get<0>(state)(_, _, *tile_iter), + get<1>(state)(_, _, smem_pipe_write.index())); + } else { + copy( + params.with(*tma_barrier, mcast_mask), + get<0>(state)(_, _, _, *tile_iter), + get<1>(state)(_, _, _, smem_pipe_write.index())); + } + + if constexpr (kAdvancePipe) + ++smem_pipe_write; + if constexpr (kAdvanceIterator) + ++tile_iter; + } + --tile_count; + } +}; + +} // namespace cutlass::fmha::collective diff --git a/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_softmax.hpp b/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_softmax.hpp new file mode 100644 index 00000000..1edfe219 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_softmax.hpp @@ -0,0 +1,378 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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" + +#include "natten/cuda/fmha_hopper/collective/fmha_common.hpp" + +namespace cutlass::fmha::collective { + +template +struct CollectiveSoftmax { + Params const& params; + CUTLASS_DEVICE CollectiveSoftmax(Params const& params) : params(params) {} + + using SumType = float; + using MaxType = ElementAccumulator; + + template + CUTLASS_DEVICE auto init( + AccPV const& acc_pv, + TiledMmaPV const& tiled_mma_pv) { + Tensor s_max = make_fragment_like( + size<0>(layout_acc_mn(tiled_mma_pv, acc_pv.layout()))); + Tensor a_sum = make_fragment_like(s_max); + return make_tuple(s_max, a_sum); + } + + CUTLASS_DEVICE float overload_exp2(float f) { + return ::exp2f(f); + } + + CUTLASS_DEVICE cutlass::half_t overload_exp2(cutlass::half_t f) { + auto a = f.raw(); + decltype(a) d; + asm("ex2.approx.f16 %0, %1;" : "=h"(d) : "h"(a)); + return cutlass::half_t::bitcast(d); + } + + CUTLASS_DEVICE float overload_max(float a, float b) { + return ::max(a, b); + } + + CUTLASS_DEVICE cutlass::half_t overload_max( + cutlass::half_t a, + cutlass::half_t b) { + return cutlass::half_t{__hmax_nan(a.to_half(), b.to_half())}; + } + + CUTLASS_DEVICE half overload_to_native(cutlass::half_t f) { + return f.to_half(); + } + + CUTLASS_DEVICE float overload_to_native(float f) { + return f; + } + + template < + class AccQK, + class TiledMmaQK, + class CountQK, + class State, + class ProblemShape> + CUTLASS_DEVICE auto step( + AccQK& acc_qk, + TiledMmaQK const& tiled_mma_qk, + CountQK const& count_qk, + State& state, + ProblemShape const& problem_shape) { + Fusion{}.before_softmax(acc_qk, count_qk, problem_shape); + Tensor acc_qk_mn = make_tensor( + acc_qk.data(), layout_acc_mn(tiled_mma_qk, acc_qk.layout())); + auto reduction_target_qk = reduction_target_n(tiled_mma_qk); + constexpr int red_rank = decltype(rank(reduction_target_qk))::value; + + auto& s_max = get<0>(state); + auto& a_sum = get<1>(state); + + // Linear reduction is faster for the first iteration + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<0>(acc_qk_mn); i++) { + s_max(i) = acc_qk_mn(i, 0); + } + CUTLASS_PRAGMA_UNROLL + for (int j = 1; j < size<1>(acc_qk_mn); j++) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<0>(acc_qk_mn); i++) { + s_max(i) = overload_max(s_max(i), acc_qk_mn(i, j)); + } + } + + for_each(make_seq{}, [&](auto r) { + CUTLASS_PRAGMA_UNROLL + for (int j = 1; j < shape(reduction_target_qk); j *= 2) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<0>(acc_qk_mn); i++) { + s_max(i) = overload_max( + s_max(i), + MaxType{__shfl_xor_sync( + uint32_t(-1), + overload_to_native(s_max(i)), + stride(reduction_target_qk) * j)}); + } + } + }); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<0>(acc_qk_mn); i++) { + MaxType local_max = s_max(i) == static_cast(-INFINITY) + ? static_cast(0) + : s_max(i); + MaxType scale = static_cast(params.scale_softmax_log2); + MaxType scale_max = scale * local_max; + + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<1>(acc_qk_mn); j++) { + acc_qk_mn(i, j) = overload_exp2(scale * acc_qk_mn(i, j) - scale_max); + } + } + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<0>(acc_qk_mn); i++) { + a_sum(i) = SumType{reduce(acc_qk_mn(i, _), cute::plus{})}; + } + } + + template < + bool kUseFusion = true, + class AccQK, + class TiledMmaQK, + class CountQK, + class State, + class AccPV, + class TiledMmaPV, + class ProblemShape> + CUTLASS_DEVICE auto step_interleave_begin( + AccQK& acc_qk, + TiledMmaQK const& tiled_mma_qk, + CountQK const& count_qk, + State& state, + AccPV& acc_pv, + TiledMmaPV const& tiled_mma_pv, + ProblemShape const& problem_shape) { + if constexpr (kUseFusion) { + Fusion{}.before_softmax(acc_qk, count_qk, problem_shape); + } + + Tensor acc_qk_mn = make_tensor( + acc_qk.data(), layout_acc_mn(tiled_mma_qk, acc_qk.layout())); + Tensor acc_pv_mn = make_tensor( + acc_pv.data(), layout_acc_mn(tiled_mma_pv, acc_pv.layout())); + + static_assert(size<0>(acc_qk_mn) == size<0>(acc_pv_mn)); + auto reduction_target_qk = reduction_target_n(tiled_mma_qk); + constexpr int red_rank = decltype(rank(reduction_target_qk))::value; + + auto& s_max = get<0>(state); + auto& a_sum = get<1>(state); + + Tensor s_max_prev = make_fragment_like(s_max); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<0>(acc_qk_mn); i++) { + s_max_prev(i) = s_max(i); + } + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<0>(acc_qk_mn); i++) { + // Linear reduction is faster here, as well + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<1>(acc_qk_mn); j++) { + s_max(i) = overload_max(s_max(i), acc_qk_mn(i, j)); + } + } + // reduce max + for_each(make_seq{}, [&](auto r) { + CUTLASS_PRAGMA_UNROLL + for (int j = 1; j < shape(reduction_target_qk); j *= 2) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<0>(acc_qk_mn); i++) { + s_max(i) = overload_max( + s_max(i), + __shfl_xor_sync( + uint32_t(-1), s_max(i), stride(reduction_target_qk) * j)); + } + } + }); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<0>(acc_pv_mn); i++) { + float s_max_cur = s_max(i) == -INFINITY ? 0.0f : s_max(i); + float scale = + ::exp2f((s_max_prev(i) - s_max_cur) * params.scale_softmax_log2); + a_sum(i) *= scale; + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<1>(acc_pv_mn); j++) { + acc_pv_mn(i, j) *= scale; + } + } + } + + template + CUTLASS_DEVICE auto step_interleave_step(AccQK_MN& acc_qk_mn, State& state) { + auto& s_max = get<0>(state); + auto& a_sum = get<1>(state); + + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<0>(acc_qk_mn); j++) { + float local_max = s_max(j) == -INFINITY ? 0.f : s_max(j); + float scale_max = params.scale_softmax_log2 * local_max; + + CUTLASS_PRAGMA_UNROLL + for (int k = 0; k < size<1>(acc_qk_mn); k++) { + acc_qk_mn(j, k) = + ::exp2f(params.scale_softmax_log2 * acc_qk_mn(j, k) - scale_max); + a_sum(j) += acc_qk_mn(j, k); + } + } + } + + template < + bool kUseFusion = true, + class AccQK, + class TiledMmaQK, + class CountQK, + class State, + class AccPV, + class TiledMmaPV, + class ProblemShape> + CUTLASS_DEVICE auto step( + AccQK& acc_qk, + TiledMmaQK const& tiled_mma_qk, + CountQK const& count_qk, + State& state, + AccPV& acc_pv, + TiledMmaPV const& tiled_mma_pv, + ProblemShape const& problem_shape) { + if constexpr (kUseFusion) { + Fusion{}.before_softmax(acc_qk, count_qk, problem_shape); + } + + Tensor acc_qk_mn = make_tensor( + acc_qk.data(), layout_acc_mn(tiled_mma_qk, acc_qk.layout())); + Tensor acc_pv_mn = make_tensor( + acc_pv.data(), layout_acc_mn(tiled_mma_pv, acc_pv.layout())); + + static_assert(size<0>(acc_qk_mn) == size<0>(acc_pv_mn)); + auto reduction_target_qk = reduction_target_n(tiled_mma_qk); + constexpr int red_rank = decltype(rank(reduction_target_qk))::value; + + auto& s_max = get<0>(state); + auto& a_sum = get<1>(state); + + Tensor s_max_prev = make_fragment_like(s_max); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<0>(acc_qk_mn); i++) { + s_max_prev(i) = s_max(i); + + // Linear reduction is faster here, as well + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<1>(acc_qk_mn); j++) { + s_max(i) = overload_max(s_max(i), acc_qk_mn(i, j)); + } + // reduce max + for_each(make_seq{}, [&](auto r) { + CUTLASS_PRAGMA_UNROLL + for (int j = 1; j < shape(reduction_target_qk); j *= 2) { + s_max(i) = overload_max( + s_max(i), + MaxType{__shfl_xor_sync( + uint32_t(-1), + overload_to_native(s_max(i)), + stride(reduction_target_qk) * j)}); + } + }); + + MaxType local_max = s_max(i) == static_cast(-INFINITY) + ? static_cast(0) + : s_max(i); + MaxType scale = static_cast(params.scale_softmax_log2); + MaxType scale_max = scale * local_max; + + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<1>(acc_qk_mn); j++) { + acc_qk_mn(i, j) = overload_exp2(scale * acc_qk_mn(i, j) - scale_max); + } + + MaxType s_max_cur = s_max(i) == static_cast(-INFINITY) + ? static_cast(0) + : s_max(i); + SumType scale_pv = overload_exp2((s_max_prev(i) - s_max_cur) * scale); + a_sum(i) *= scale_pv; + + using ElementPV = typename AccPV::value_type; + ElementPV scale_pv_ele = static_cast(scale_pv); + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<1>(acc_pv_mn); j++) { + acc_pv_mn(i, j) *= scale_pv_ele; + } + a_sum(i) += SumType{reduce(acc_qk_mn(i, _), cute::plus{})}; + } + } + + template + CUTLASS_DEVICE auto tail( + State& state, + AccPV& acc_pv, + TiledMmaPV const& tiled_mma_pv) { + auto& s_max = get<0>(state); + auto& a_sum = get<1>(state); + + Tensor acc_pv_mn = make_tensor( + acc_pv.data(), layout_acc_mn(tiled_mma_pv, acc_pv.layout())); + + auto reduction_target = reduction_target_n(tiled_mma_pv); + constexpr int red_rank = decltype(rank(reduction_target))::value; + for_each(make_seq{}, [&](auto r) { + CUTLASS_PRAGMA_UNROLL + for (int j = 1; j < shape(reduction_target); j *= 2) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<0>(acc_pv_mn); i++) { + a_sum(i) = + a_sum(i) + + __shfl_xor_sync( + uint32_t(-1), a_sum(i), stride(reduction_target) * j); + } + } + }); + + Tensor acc_mn = make_tensor( + acc_pv.data(), layout_acc_mn(tiled_mma_pv, acc_pv.layout())); + + Tensor lse = make_fragment_like(a_sum); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<0>(acc_mn); i++) { + float sum = a_sum(i); + float inv_sum = (sum == 0.f || sum != sum) ? 1.f : __frcp_rn(sum); + lse(i) = (sum == 0.f || sum != sum) + ? INFINITY + : s_max(i) * params.scale_softmax + __logf(sum); + float scale = params.rp_dropout * inv_sum; + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<1>(acc_mn); j++) { + acc_mn(i, j) *= scale; + } + } + + return lse; + } +}; + +} // namespace cutlass::fmha::collective diff --git a/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_tma.hpp b/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_tma.hpp new file mode 100644 index 00000000..03f8a7ba --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_tma.hpp @@ -0,0 +1,715 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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/gemm/collective/collective_builder.hpp" + +#include "natten/cuda/fmha_hopper/collective/fmha_collective_load.hpp" +#include "natten/cuda/fmha_hopper/collective/fmha_collective_softmax.hpp" +#include "natten/cuda/fmha_hopper/collective/fmha_common.hpp" +#include "natten/cuda/fmha_hopper/collective/fmha_varlen.hpp" +#include "natten/cuda/fmha_hopper/kernel/fmha_options.hpp" + +namespace cutlass::fmha::collective { + +using namespace cute; +using cutlass::fmha::kernel::find_option_t; +using cutlass::fmha::kernel::Tag; + +template < + typename Element_, + typename ElementAccumulator_, + typename TileShape_, // BlockQO, BlockKV, BlockHead + class Fusion, + class... Options> +struct FmhaMainloopTmaSm90 { + using Element = Element_; + using ElementAccumulator = ElementAccumulator_; + using TileShape = TileShape_; + + // Options + using kClusterM = find_option_t, Options...>; + static constexpr int StageCount = + find_option_t, Options...>::value; + static constexpr int StageCountQ = + find_option_t, Options...>::value; + + using StagesQ = cutlass::gemm::collective::StageCount; + using Stages = cutlass::gemm::collective::StageCount; + using ClusterShape = Shape; + + // 16B alignment lets us use TMA + static constexpr int Alignment = 16 / sizeof(Element); + + using TileShapeQK = TileShape; + using TileShapePV = decltype(select<0, 2, 1>(TileShapeQK{})); + + using LayoutQKV = cute::tuple>; + using LayoutQ = LayoutQKV; + using LayoutK = LayoutQKV; + using LayoutV = LayoutQKV; + + using CollectiveMmaQK = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm90, + cutlass::arch::OpClassTensorOp, + Element, + LayoutQ, + Alignment, + Element, + LayoutK, + Alignment, + ElementAccumulator, + TileShapeQK, + ClusterShape, + Stages, + cutlass::gemm::KernelTmaWarpSpecialized>::CollectiveOp; + + using CollectiveMmaPV = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm90, + cutlass::arch::OpClassTensorOp, + // the stride for A does not matter since we do not load from smem at all + Element, + LayoutK, + Alignment, + Element, + decltype(select<1, 0, 2>(LayoutV{})), + Alignment, + ElementAccumulator, + TileShapePV, + ClusterShape, + Stages, + cutlass::gemm::KernelTmaWarpSpecialized>::CollectiveOp; + + using TiledMmaQK = typename CollectiveMmaQK::TiledMma; + using TiledMmaPV = + decltype(convert_to_gmma_rs(typename CollectiveMmaPV::TiledMma{})); + + using SmemLayoutQ = decltype(unstageSmemLayout( + typename CollectiveMmaQK::SmemLayoutA{}, + Int{})); + using SmemLayoutK = typename CollectiveMmaQK::SmemLayoutB; + using SmemLayoutV = typename CollectiveMmaPV::SmemLayoutB; + + using MainloopPipeline = cutlass::PipelineTmaAsync; + using MainloopPipelineQ = cutlass::PipelineTmaAsync; + + using PipelineState = + typename cutlass::PipelineState; + using PipelineStateQ = + typename cutlass::PipelineState; + + using TileShapeOut = TileShapePV; + using TiledMmaOut = TiledMmaPV; + using ElementOut = ElementAccumulator; + + struct SharedStorage { + cute::array_aligned> smem_q; + union { + cute::array_aligned> smem_k; + cute::array_aligned> smem_v; + }; + }; + + struct Arguments { + const Element* ptr_Q; + LayoutQ dQ; + const Element* ptr_K; + LayoutK dK; + const Element* ptr_V; + LayoutV dV; + + // if zero, defaults to 1/sqrt(D) + float scale_softmax = 0.0f; + }; + + 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; + TMA_K tma_load_k; + TMA_V tma_load_v; + + float scale_softmax; + float scale_softmax_log2; + float rp_dropout; + }; + + using LoadQ = cutlass::fmha::collective::CollectiveLoadTma< + cutlass::fmha::collective::LoadKind::kQ, + MainloopPipelineQ, + Element, + SmemLayoutQ, + TMA_Q>; + + using LoadK = cutlass::fmha::collective::CollectiveLoadTma< + cutlass::fmha::collective::LoadKind::kK, + MainloopPipeline, + Element, + SmemLayoutK, + TMA_K>; + + using LoadV = cutlass::fmha::collective::CollectiveLoadTma< + cutlass::fmha::collective::LoadKind::kV, + MainloopPipeline, + Element, + SmemLayoutV, + TMA_V>; + + static_assert( + size(typename CollectiveMmaQK::TiledMma{}) == + size(typename CollectiveMmaPV::TiledMma{})); + + static const int MaxThreadsPerBlock = + size(typename CollectiveMmaQK::TiledMma{}); + + template + static bool can_implement( + ProblemShape const& problem_size, + Arguments const& args) { + // NOTE (ali): the original Hopper FMHA kernel does an alignment check on + // seqlen_Q as well, citing TMA requirement. However, it does not look like + // TMA runs into any issues, because cases that violate this run perfectly + // fine. + // + // Backprop requires this extra check though, and the point of failure is + // LSE. + // + // This does not cause any issues with FNA kernels, since they already pad + // for tile size, and tile sizes are usually powers of two, and at least 64, + // so divisibility by the required alignment 16 is a guarantee, though not + // explicitly asserted. + // + // Adding in padding along seqlen_Q requires modifications in the bwd mask + // as well. + // + // Therefore, we only pad LSE, and not any other tensors to get around this + // issue in bwd, instead of padding Q, dQ, O and dO as well. + // + // If we pad those, we would have to modify the residual mask as well. + // + return (get<4>(problem_size) <= get<2>(TileShape{})) && + ((get<4>(problem_size) % Alignment) == 0) /* && + ((get<2>(problem_size) % Alignment) == 0)*/ + ; + } + + template + static Params to_underlying_arguments( + ProblemShape const& problem_size, + Arguments const& args, + void* workspace) { + auto problem_shape_qk = apply_variable_length_scheduler(make_shape( + get<2>(problem_size), + get<3>(problem_size), + get<4>(problem_size), + make_shape(get<0>(problem_size), get<1>(problem_size)))); + auto params_qk = CollectiveMmaQK::to_underlying_arguments( + problem_shape_qk, + typename CollectiveMmaQK::Arguments{ + args.ptr_Q, + args.dQ, + args.ptr_K, + args.dK, + }, + /*workspace=*/nullptr); + + auto problem_shape_pv = select<0, 2, 1, 3>(problem_shape_qk); + auto params_pv = CollectiveMmaPV::to_underlying_arguments( + problem_shape_pv, + typename CollectiveMmaPV::Arguments{ + args.ptr_K, + args.dK, // never used, dummy + args.ptr_V, + select<1, 0, 2>(args.dV), + }, + /*workspace=*/nullptr); + + float scale_softmax = args.scale_softmax; + if (scale_softmax == 0.0f) { + scale_softmax = 1.0f / (float)std::sqrt(get<4>(problem_size)); + } + float log2_e = static_cast(std::log2(std::exp(1.0))); + + return Params{ + params_qk.tma_load_a, + params_qk.tma_load_b, + params_pv.tma_load_b, + scale_softmax, + log2_e * scale_softmax, + 1.0f}; + } + + 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 auto compute( + int block_rank_in_cluster, + BlkCoord const& blk_coord, + Params const& params, + ProblemShape const& problem_size, + MainloopPipeline& pipeline, + PipelineState& smem_pipe_read, + PipelineState& smem_pipe_write, + MainloopPipelineQ& pipeline_q, + PipelineStateQ& smem_pipe_read_q, + PipelineStateQ& smem_pipe_write_q, + SharedStorage& storage) { + int warp_idx = cutlass::canonical_warp_idx_sync(); + int thread_idx = threadIdx.x; + + PipelineState smem_pipe_release = smem_pipe_read; + [[maybe_unused]] PipelineStateQ smem_pipe_release_q = smem_pipe_read_q; + + int fusion_tile_count = 0; + if constexpr (is_problem_shape_variable_length(ProblemShape{})) { + auto problem_size_ = + apply_variable_length(problem_size, get<2, 0>(blk_coord)); + fusion_tile_count = + Fusion{}.get_trip_count(blk_coord, TileShape{}, problem_size_); + } else { + fusion_tile_count = + Fusion{}.get_trip_count(blk_coord, TileShape{}, problem_size); + } + + LoadQ load_q{params.tma_load_q, pipeline_q, storage.smem_q}; + auto load_state_q = + load_q.init_state(_0{}, problem_size, TileShapeQK{}, blk_coord, 1); + + LoadK load_k{params.tma_load_k, pipeline, storage.smem_k}; + auto load_state_k = load_k.init_state( + block_rank_in_cluster, + problem_size, + TileShapeQK{}, + blk_coord, + fusion_tile_count); + + LoadV load_v{params.tma_load_v, pipeline, storage.smem_v}; + auto load_state_v = load_v.init_state( + block_rank_in_cluster, + problem_size, + TileShapePV{}, + blk_coord, + fusion_tile_count); + + // Set predicate for the lowest lane_id in the warp + int lane_predicate = cute::elect_one_sync(); + + // Issue TmaLoads (Prologue fetches) + if (warp_idx == 0) { + auto q_tile_iter = cute::make_coord_iterator(1); + int q_tile_count = 1; + load_q.step( + q_tile_iter, + load_state_q, + smem_pipe_write_q, + lane_predicate, + q_tile_count); + } + + // Loop over K elems + auto k_tile_iter = cute::make_coord_iterator(fusion_tile_count); + + int k_tile_count_tma = 2 * fusion_tile_count; + + uint16_t mcast_mask_b = 0; + + if (warp_idx == 0 && lane_predicate == 1) { + if constexpr (cute::is_same_v< + typename CollectiveMmaQK::GmemTiledCopyB, + SM90_TMA_LOAD_MULTICAST>) { + auto block_layout = Layout{}; // (m,n) -> block_id + for (int m = 0; m < size<0>(block_layout); ++m) { + mcast_mask_b |= (uint16_t(1) << block_layout(m, _0{}, Int<0>{})); + } + } + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < StageCount; i++) { + if (i % 2 == 0) { + load_k.template step( + k_tile_iter, + load_state_k, + smem_pipe_write, + lane_predicate, + k_tile_count_tma, + mcast_mask_b); + } else { + load_v.template step( + k_tile_iter, + load_state_k, + smem_pipe_write, + lane_predicate, + k_tile_count_tma, + mcast_mask_b); + } + } + } + + TiledMmaQK tiled_mma_qk; + auto thr_mma_qk = tiled_mma_qk.get_thread_slice(thread_idx); + + // Mainloop setup QK + 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 tSsQ = thr_mma_qk.partition_A(sQ); // (MMA,MMA_M,MMA_K,PIPE) + Tensor tSsK = thr_mma_qk.partition_B(sK); // (MMA,MMA_N,MMA_K,PIPE) + Tensor tSrQ = thr_mma_qk.make_fragment_A(tSsQ); // (MMA,MMA_N,MMA_K,PIPE) + Tensor tSrK = thr_mma_qk.make_fragment_B(tSsK); // (MMA,MMA_M,MMA_N,PIPE) + + // Prepare: MMA PV + TiledMmaPV tiled_mma_pv; + auto thr_mma_pv = tiled_mma_pv.get_thread_slice(thread_idx); + + // Mainloop setup PV + Tensor sV = + make_tensor(make_smem_ptr(storage.smem_v.data()), SmemLayoutV{}); + + Tensor tOsV = thr_mma_pv.partition_B(sV); // (MMA,MMA_N,MMA_K,PIPE) + Tensor tOrV = thr_mma_pv.make_fragment_B(tOsV); // (MMA,MMA_M,MMA_N,PIPE) + + auto problem_size_f = + apply_variable_length(problem_size, get<2, 0>(blk_coord)); + int k_tile_count = Fusion{}.get_unmasked_trip_count( + blk_coord, TileShape{}, problem_size_f); + + pipeline_q.consumer_wait(smem_pipe_read_q); + + // mapping into QK accumulator + Tensor cP = make_identity_tensor(take<0, 2>(TileShapeQK{})); + Tensor tPcP = thr_mma_qk.partition_C(cP); + int m_block = get<0>(blk_coord); + tPcP.data() = tPcP.data() + E<0>{} * m_block * get<0>(TileShapeQK{}); + + // Allocate PV acc + Tensor acc_pv = + partition_fragment_C(tiled_mma_pv, take<0, 2>(TileShapePV{})); + + cutlass::fmha::collective:: + CollectiveSoftmax + softmax{params}; + auto softmax_state = softmax.init(acc_pv, tiled_mma_pv); + + if (true) { + --k_tile_count; + // Allocate QK acc + Tensor acc_qk = + partition_fragment_C(tiled_mma_qk, take<0, 2>(TileShapeQK{})); + + pipeline.consumer_wait(smem_pipe_read); + + // MMA QK + warpgroup_fence_operand(acc_qk); + warpgroup_arrive(); + + gemm_zero_acc( + tiled_mma_qk, + tSrQ(_, _, _, _0{}), + tSrK(_, _, _, smem_pipe_read.index()), + acc_qk); + warpgroup_commit_batch(); + + ++smem_pipe_read; + + // Wait for the pipeline MMAs to drain + warpgroup_wait<0>(); + warpgroup_fence_operand(acc_qk); + + softmax.step(acc_qk, tiled_mma_qk, tPcP, softmax_state, problem_size_f); + + Tensor acc_qk_fixed = + make_fragment_like(convert_c_layout_to_a_layout( + acc_qk.layout(), + shape<1>(typename decltype(tiled_mma_pv)::LayoutA_TV{}))); + + Tensor acc_qk_input = make_tensor(acc_qk_fixed.data(), acc_qk.layout()); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(acc_qk); i++) { + acc_qk_input(i) = static_cast(acc_qk(i)); + } + + pipeline.consumer_wait(smem_pipe_read); + + // MMA PV + warpgroup_fence_operand(acc_pv); + warpgroup_fence_operand(acc_qk_fixed); + warpgroup_arrive(); + + gemm_zero_acc( + tiled_mma_pv, + acc_qk_fixed, + tOrV(_, _, _, smem_pipe_read.index()), + acc_pv); + warpgroup_commit_batch(); + + // + // Advance the pipe + // + + // Advance consumer pipeline + ++smem_pipe_read; + + pipeline.consumer_release(smem_pipe_release); + ++smem_pipe_release; + + tPcP.data() = tPcP.data() + E<1>{} * get<1>(TileShapeQK{}); + } + + CUTLASS_PRAGMA_NO_UNROLL + for (; k_tile_count > 0; --k_tile_count) { + // Allocate QK acc + Tensor acc_qk = + partition_fragment_C(tiled_mma_qk, take<0, 2>(TileShapeQK{})); + + pipeline.consumer_wait(smem_pipe_read); + + // MMA QK + warpgroup_fence_operand(acc_qk); + warpgroup_arrive(); + + gemm_zero_acc( + tiled_mma_qk, + tSrQ(_, _, _, _0{}), + tSrK(_, _, _, smem_pipe_read.index()), + acc_qk); + warpgroup_commit_batch(); + + ++smem_pipe_read; + + if (warp_idx == 0) { + load_k.template step( + k_tile_iter, + load_state_k, + smem_pipe_write, + lane_predicate, + k_tile_count_tma, + mcast_mask_b); + } + + // Wait for the pipeline MMAs to drain + warpgroup_wait<0>(); + warpgroup_fence_operand(acc_qk); + warpgroup_fence_operand(acc_pv); + + softmax.template step_interleave_begin( + acc_qk, + tiled_mma_qk, + tPcP, + softmax_state, + acc_pv, + tiled_mma_pv, + problem_size_f); + + pipeline.consumer_release(smem_pipe_release); + + ++smem_pipe_release; + + pipeline.consumer_wait(smem_pipe_read); + + // MMA PV + auto layout_qk_input = convert_c_layout_to_a_layout( + acc_qk.layout(), + shape<1>(typename decltype(tiled_mma_pv)::LayoutA_TV{})); + + Tensor acc_qk_input = make_tensor(acc_qk.data(), layout_qk_input); + + static_assert(decltype(size<1>(layout_qk_input) == _1{})::value); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<2>(tOrV); i++) { + Tensor acc_qk_element = + make_fragment_like(layout_qk_input(_, _0{}, _0{})); + Tensor acc_qk_element_mk = tensor_op_mk_v(tiled_mma_pv, acc_qk_element); + Tensor acc_qk_input_mk = + tensor_op_mk_v(tiled_mma_pv, acc_qk_input(_, _0{}, i)); + softmax.step_interleave_step(acc_qk_input_mk, softmax_state); + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size(acc_qk_element_mk); j++) { + acc_qk_element_mk(j) = static_cast(acc_qk_input_mk(j)); + } + warpgroup_arrive(); + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<1>(tOrV); j++) { + cute::gemm( + tiled_mma_pv, + acc_qk_element, + tOrV(_, j, i, smem_pipe_read.index()), + acc_pv(_, _0{}, j)); + } + } + warpgroup_commit_batch(); + + // Wait for the pipeline MMAs to drain + pipeline.consumer_release(smem_pipe_release); + ++smem_pipe_release; + + ++smem_pipe_read; + + if (warp_idx == 0) { + load_v.template step( + k_tile_iter, + load_state_v, + smem_pipe_write, + lane_predicate, + k_tile_count_tma, + mcast_mask_b); + } + + tPcP.data() = tPcP.data() + E<1>{} * get<1>(TileShapeQK{}); + } + + k_tile_count += + Fusion{}.get_masked_trip_count(blk_coord, TileShape{}, problem_size_f); + + CUTLASS_PRAGMA_NO_UNROLL + for (; k_tile_count > 0; --k_tile_count) { + // Allocate QK acc + Tensor acc_qk = + partition_fragment_C(tiled_mma_qk, take<0, 2>(TileShapeQK{})); + + pipeline.consumer_wait(smem_pipe_read); + + // MMA QK + warpgroup_fence_operand(acc_qk); + warpgroup_arrive(); + + gemm_zero_acc( + tiled_mma_qk, + tSrQ(_, _, _, _0{}), + tSrK(_, _, _, smem_pipe_read.index()), + acc_qk); + warpgroup_commit_batch(); + + ++smem_pipe_read; + + if (warp_idx == 0) { + load_k.template step( + k_tile_iter, + load_state_k, + smem_pipe_write, + lane_predicate, + k_tile_count_tma, + mcast_mask_b); + } + + // Wait for the pipeline MMAs to drain + warpgroup_wait<0>(); + warpgroup_fence_operand(acc_qk); + warpgroup_fence_operand(acc_pv); + + softmax.step_interleave_begin( + acc_qk, + tiled_mma_qk, + tPcP, + softmax_state, + acc_pv, + tiled_mma_pv, + problem_size_f); + + pipeline.consumer_release(smem_pipe_release); + + ++smem_pipe_release; + + pipeline.consumer_wait(smem_pipe_read); + + // MMA PV + auto layout_qk_input = convert_c_layout_to_a_layout( + acc_qk.layout(), + shape<1>(typename decltype(tiled_mma_pv)::LayoutA_TV{})); + + Tensor acc_qk_input = make_tensor(acc_qk.data(), layout_qk_input); + + static_assert(decltype(size<1>(layout_qk_input) == _1{})::value); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<2>(tOrV); i++) { + Tensor acc_qk_element = + make_fragment_like(layout_qk_input(_, _0{}, _0{})); + Tensor acc_qk_element_mk = tensor_op_mk_v(tiled_mma_pv, acc_qk_element); + Tensor acc_qk_input_mk = + tensor_op_mk_v(tiled_mma_pv, acc_qk_input(_, _0{}, i)); + softmax.step_interleave_step(acc_qk_input_mk, softmax_state); + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size(acc_qk_element_mk); j++) { + acc_qk_element_mk(j) = static_cast(acc_qk_input_mk(j)); + } + warpgroup_arrive(); + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<1>(tOrV); j++) { + cute::gemm( + tiled_mma_pv, + acc_qk_element, + tOrV(_, j, i, smem_pipe_read.index()), + acc_pv(_, _0{}, j)); + } + } + warpgroup_commit_batch(); + + // Wait for the pipeline MMAs to drain + pipeline.consumer_release(smem_pipe_release); + ++smem_pipe_release; + + ++smem_pipe_read; + + if (warp_idx == 0) { + load_v.template step( + k_tile_iter, + load_state_v, + smem_pipe_write, + lane_predicate, + k_tile_count_tma, + mcast_mask_b); + } + + tPcP.data() = tPcP.data() + E<1>{} * get<1>(TileShapeQK{}); + } + + // Wait for the pipeline MMAs to drain + warpgroup_wait<0>(); + warpgroup_fence_operand(acc_pv); + + Tensor lse = softmax.tail(softmax_state, acc_pv, tiled_mma_pv); + + return make_tuple(acc_pv, lse); + } +}; + +} // namespace cutlass::fmha::collective diff --git a/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_tma_warpspecialized.hpp b/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_tma_warpspecialized.hpp new file mode 100644 index 00000000..0b1df1d2 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_tma_warpspecialized.hpp @@ -0,0 +1,775 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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/gemm/collective/collective_builder.hpp" + +#include "natten/cuda/fmha_hopper/collective/fmha_collective_load.hpp" +#include "natten/cuda/fmha_hopper/collective/fmha_collective_softmax.hpp" +#include "natten/cuda/fmha_hopper/collective/fmha_common.hpp" +#include "natten/cuda/fmha_hopper/collective/fmha_varlen.hpp" +#include "natten/cuda/fmha_hopper/kernel/fmha_options.hpp" + +namespace cutlass::fmha::collective { + +using namespace cute; +using cutlass::fmha::kernel::find_option_t; +using cutlass::fmha::kernel::Tag; + +template < + class Element_, + class ElementAccumulatorQK_, + class ElementAccumulatorPV_, + class TileShape_, // SeqQ, SeqKV, Head + class LayoutQ_, + class LayoutK_, + class LayoutV_, // SeqX, Head, (Batches) + class Fusion, + class... Options> +struct FmhaMainloopTmaWarpSpecializedSm90 { + using Element = Element_; + using ElementAccumulatorQK = ElementAccumulatorQK_; + using ElementAccumulatorPV = ElementAccumulatorPV_; + using TileShape = TileShape_; + + using LayoutQ = LayoutQ_; + using LayoutK = LayoutK_; + using LayoutV = LayoutV_; + + // Options + static constexpr bool kIsPersistent = + find_option_t::value; + static constexpr bool kIsMainloopLocked = + find_option_t::value; + + static constexpr int NumLoadWarpGroups = 1; + static constexpr int NumMmaWarpGroups = + find_option_t, Options...>::value; + static constexpr int StageCount = + find_option_t, Options...>::value; + static constexpr int StageCountQ = + find_option_t, Options...>::value; + + static const int kOuterLoads = 1; + using StagesQ = cutlass::gemm::collective::StageCount; + using Stages = cutlass::gemm::collective::StageCount; + using ClusterShape = Shape<_1, _1, _1>; + static_assert(StagesQ::value >= NumMmaWarpGroups); + static_assert(Stages::value >= 2); + + // 16B alignment lets us use TMA + static constexpr int Alignment = 16 / sizeof(Element); + + using TileShapeQK = Shape< + decltype(tuple_element_t<0, TileShape>{} / Int{}), + tuple_element_t<1, TileShape>, + tuple_element_t<2, TileShape>>; + + using TileShapePV = decltype(select<0, 2, 1>(TileShapeQK{})); + + using CollectiveMmaQK = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm90, + cutlass::arch::OpClassTensorOp, + Element, + LayoutQ, + Alignment, + Element, + LayoutK, + Alignment, + ElementAccumulatorQK, + TileShapeQK, + ClusterShape, + Stages, + cutlass::gemm::KernelTmaWarpSpecialized>::CollectiveOp; + + using CollectiveMmaPV = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm90, + cutlass::arch::OpClassTensorOp, + // the stride for A does not matter since we do not load from smem at all + Element, + LayoutK, + Alignment, + Element, + decltype(select<1, 0, 2>(LayoutV{})), + Alignment, + ElementAccumulatorPV, + TileShapePV, + ClusterShape, + Stages, + cutlass::gemm::KernelTmaWarpSpecialized>::CollectiveOp; + + using TiledMmaQK = typename CollectiveMmaQK::TiledMma; + using TiledMmaPV = + decltype(convert_to_gmma_rs(typename CollectiveMmaPV::TiledMma{})); + + using SmemLayoutQ = decltype(unstageSmemLayout( + typename CollectiveMmaQK::SmemLayoutA{}, + Int{})); + using SmemLayoutK = typename CollectiveMmaQK::SmemLayoutB; + using SmemLayoutV = typename CollectiveMmaPV::SmemLayoutB; + + using MainloopPipeline = cutlass::PipelineTmaAsync; + using MainloopPipelineQ = cutlass::PipelineTmaAsync; + + using PipelineState = + typename cutlass::PipelineState; + using PipelineStateQ = + typename cutlass::PipelineState; + + static constexpr int kInnerLoadBytes = + size(SmemLayoutK{}(_, _, _0{})) * sizeof(Element); + static constexpr int kOuterLoadBytes = + size(SmemLayoutQ{}(_, _, _0{})) * sizeof(Element); + + using TileShapeOut = TileShapePV; + using TiledMmaOut = TiledMmaPV; + using ElementOut = ElementAccumulatorPV; + + struct SharedStorage { + cute::array_aligned> smem_q; + union { + cute::array_aligned> smem_k; + cute::array_aligned> smem_v; + }; + }; + + struct Arguments { + const Element* ptr_Q; + LayoutQ dQ; + const Element* ptr_K; + LayoutK dK; + const Element* ptr_V; + LayoutV dV; + + // if zero, defaults to 1/sqrt(D) + float scale_softmax = 0.0f; + }; + + 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; + TMA_K tma_load_k; + TMA_V tma_load_v; + + float scale_softmax; + float scale_softmax_log2; + float rp_dropout; + }; + + using LoadQ = cutlass::fmha::collective::CollectiveLoadTma< + cutlass::fmha::collective::LoadKind::kQ, + MainloopPipelineQ, + Element, + SmemLayoutQ, + TMA_Q>; + + using LoadK = cutlass::fmha::collective::CollectiveLoadTma< + cutlass::fmha::collective::LoadKind::kK, + MainloopPipeline, + Element, + SmemLayoutK, + TMA_K>; + + using LoadV = cutlass::fmha::collective::CollectiveLoadTma< + cutlass::fmha::collective::LoadKind::kV, + MainloopPipeline, + Element, + SmemLayoutV, + TMA_V>; + + static_assert( + size(typename CollectiveMmaQK::TiledMma{}) == + size(typename CollectiveMmaPV::TiledMma{})); + + template + static bool can_implement( + ProblemShape const& problem_size, + Arguments const& args) { + // NOTE (ali): the original Hopper FMHA kernel does an alignment check on + // seqlen_Q as well, citing TMA requirement. However, it does not look like + // TMA runs into any issues, because cases that violate this run perfectly + // fine. + // + // Backprop requires this extra check though, and the point of failure is + // LSE. + // + // This does not cause any issues with FNA kernels, since they already pad + // for tile size, and tile sizes are usually powers of two, and at least 64, + // so divisibility by the required alignment 16 is a guarantee, though not + // explicitly asserted. + // + // Adding in padding along seqlen_Q requires modifications in the bwd mask + // as well. + // + // Therefore, we only pad LSE, and not any other tensors to get around this + // issue in bwd, instead of padding Q, dQ, O and dO as well. + // + // If we pad those, we would have to modify the residual mask as well. + // + return true && (get<4>(problem_size) <= get<2>(TileShape{})) && + ((get<4>(problem_size) % Alignment) == 0) /* && + ((get<2>(problem_size) % Alignment) == 0)*/ + ; + } + + template + static Params to_underlying_arguments( + ProblemShape const& problem_size, + Arguments const& args, + void* workspace) { + auto problem_shape_qk = apply_variable_length_scheduler(make_shape( + get<2>(problem_size), + get<3>(problem_size), + get<4>(problem_size), + make_shape(get<0>(problem_size), get<1>(problem_size)))); + auto params_qk = CollectiveMmaQK::to_underlying_arguments( + problem_shape_qk, + typename CollectiveMmaQK::Arguments{ + args.ptr_Q, + args.dQ, + args.ptr_K, + args.dK, + }, + /*workspace=*/nullptr); + + auto problem_shape_pv = select<0, 2, 1, 3>(problem_shape_qk); + auto params_pv = CollectiveMmaPV::to_underlying_arguments( + problem_shape_pv, + typename CollectiveMmaPV::Arguments{ + args.ptr_K, + args.dK, // never used, dummy + args.ptr_V, + select<1, 0, 2>(args.dV), + }, + /*workspace=*/nullptr); + + float scale_softmax = args.scale_softmax; + if (scale_softmax == 0.0f) { + scale_softmax = 1.0f / (float)std::sqrt(get<4>(problem_size)); + } + float log2_e = static_cast(std::log2(std::exp(1.0))); + + return Params{ + params_qk.tma_load_a, + params_qk.tma_load_b, + params_pv.tma_load_b, + scale_softmax, + log2_e * scale_softmax, + 1.0f}; + } + + 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 < + bool kLoadQ, + class BlkCoord, + class ProblemShape, + class LoadWarpBarrier> + CUTLASS_DEVICE void load_kv_maybe_q( + int block_rank_in_cluster, + BlkCoord const& blk_coord, + Params const& params, + ProblemShape const& problem_size, + MainloopPipeline& pipeline, + PipelineState& smem_pipe_write, + MainloopPipelineQ& pipeline_q, + PipelineStateQ& smem_pipe_write_q, + SharedStorage& storage, + LoadWarpBarrier& load_warp_barrier, + bool do_barrier) { + int fusion_tile_count = 0; + if constexpr (is_problem_shape_variable_length(ProblemShape{})) { + auto problem_size_ = + apply_variable_length(problem_size, get<2, 0>(blk_coord)); + fusion_tile_count = + Fusion{}.get_trip_count(blk_coord, TileShape{}, problem_size_); + } else { + fusion_tile_count = + Fusion{}.get_trip_count(blk_coord, TileShape{}, problem_size); + } + + int lane_predicate = cute::elect_one_sync(); + + uint16_t mcast_mask_b = 0; + + if (lane_predicate == 1) { + if constexpr (cute::is_same_v< + typename CollectiveMmaQK::GmemTiledCopyB, + SM90_TMA_LOAD_MULTICAST>) { + auto block_layout = Layout{}; // (m,n) -> block_id + for (int m = 0; m < size<0>(block_layout); ++m) { + mcast_mask_b |= (uint16_t(1) << block_layout(m, _0{}, Int<0>{})); + } + } + } + + auto q_tile_iter = cute::make_coord_iterator(Int{}); + [[maybe_unused]] int q_tile_count = NumMmaWarpGroups; + + auto k_tile_iter = cute::make_coord_iterator(fusion_tile_count); + int k_tile_count = 2 * fusion_tile_count; + + LoadQ load_q{params.tma_load_q, pipeline_q, storage.smem_q}; + auto load_state_q = load_q.init_state( + _0{}, problem_size, TileShapeQK{}, blk_coord, NumMmaWarpGroups); + + LoadK load_k{params.tma_load_k, pipeline, storage.smem_k}; + auto load_state_k = load_k.init_state( + block_rank_in_cluster, + problem_size, + TileShapeQK{}, + blk_coord, + fusion_tile_count); + + LoadV load_v{params.tma_load_v, pipeline, storage.smem_v}; + auto load_state_v = load_v.init_state( + block_rank_in_cluster, + problem_size, + TileShapePV{}, + blk_coord, + fusion_tile_count); + + if constexpr (kLoadQ) { + load_q.step( + q_tile_iter, + load_state_q, + smem_pipe_write_q, + lane_predicate, + q_tile_count); + } + + load_k.template step( + k_tile_iter, + load_state_k, + smem_pipe_write, + lane_predicate, + k_tile_count, + mcast_mask_b); + + if constexpr (kLoadQ) { + load_q.step( + q_tile_iter, + load_state_q, + smem_pipe_write_q, + lane_predicate, + q_tile_count); + } + + if constexpr (!kLoadQ) { + if (do_barrier) { + load_warp_barrier.arrive(); + load_warp_barrier.wait(/*phase=*/0); + do_barrier = false; + } + } + + load_v.template step( + k_tile_iter, + load_state_v, + smem_pipe_write, + lane_predicate, + k_tile_count, + mcast_mask_b); + + if constexpr (kLoadQ) { + while (q_tile_count > 0) { + load_q.step( + q_tile_iter, + load_state_q, + smem_pipe_write_q, + lane_predicate, + q_tile_count); + } + } + + CUTLASS_PRAGMA_NO_UNROLL + while (k_tile_count > 0) { + load_k.template step( + k_tile_iter, + load_state_k, + smem_pipe_write, + lane_predicate, + k_tile_count, + mcast_mask_b); + load_v.template step( + k_tile_iter, + load_state_v, + smem_pipe_write, + lane_predicate, + k_tile_count, + mcast_mask_b); + } + } + + template + CUTLASS_DEVICE void load_maybe_q( + BlkCoord const& blk_coord, + Params const& params, + ProblemShape const& problem_size, + MainloopPipelineQ& pipeline_q, + PipelineStateQ& smem_pipe_write_q, + SharedStorage& storage, + LoadWarpBarrier& load_warp_barrier, + bool do_barrier) { + int lane_predicate = cute::elect_one_sync(); + + LoadQ load_q{params.tma_load_q, pipeline_q, storage.smem_q}; + auto load_state_q = load_q.init_state( + _0{}, problem_size, TileShapeQK{}, blk_coord, NumMmaWarpGroups); + + auto q_tile_iter = cute::make_coord_iterator(Int{}); + + CUTLASS_PRAGMA_UNROLL + for (int q_tile_count = 0; q_tile_count < NumMmaWarpGroups; + q_tile_count++) { + int count = 1; + load_q.step( + q_tile_iter, load_state_q, smem_pipe_write_q, lane_predicate, count); + if (q_tile_count == 0 && do_barrier) { + load_warp_barrier.arrive(); + load_warp_barrier.wait(/*phase=*/0); + do_barrier = false; + } + } + } + + template < + class BlkCoord, + class ProblemShape, + class MainloopPipelineReducer, + class PipelineStateReducer> + CUTLASS_DEVICE void reduce( + BlkCoord const& blk_coord, + Params const& params, + ProblemShape const& problem_size, + MainloopPipelineReducer& pipeline_reducer, + PipelineStateReducer& smem_pipe_write_reducer, + SharedStorage& storage) { /* no-op */ } + + template < + class BlkCoord, + class ProblemShape, + class MainloopPipelineReducer, + class PipelineStateReducer, + class MathWgOrderBarrier> + CUTLASS_DEVICE auto compute( + BlkCoord const& blk_coord, + BlkCoord const& wg_coord, + Params const& params, + ProblemShape const& problem_size, + MainloopPipeline& pipeline, + PipelineState& smem_pipe_read, + MainloopPipelineQ& pipeline_q, + PipelineStateQ& smem_pipe_read_q, + MainloopPipelineReducer&, + PipelineStateReducer&, + SharedStorage& storage, + MathWgOrderBarrier& math_wg_order_barrier) { + int thread_idx = int(threadIdx.x); + + PipelineState smem_pipe_release = smem_pipe_read; + PipelineStateQ smem_pipe_release_q = smem_pipe_read_q; + + TiledMmaQK tiled_mma_qk; + auto thr_mma_qk = tiled_mma_qk.get_thread_slice(thread_idx); + + // Mainloop setup QK + 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 tSsQ = thr_mma_qk.partition_A(sQ); // (MMA,MMA_M,MMA_K,PIPE) + Tensor tSsK = thr_mma_qk.partition_B(sK); // (MMA,MMA_N,MMA_K,PIPE) + Tensor tSrQ = thr_mma_qk.make_fragment_A(tSsQ); // (MMA,MMA_N,MMA_K,PIPE) + Tensor tSrK = thr_mma_qk.make_fragment_B(tSsK); // (MMA,MMA_M,MMA_N,PIPE) + + // Prepare: MMA PV + TiledMmaPV tiled_mma_pv; + auto thr_mma_pv = tiled_mma_pv.get_thread_slice(thread_idx); + + // Mainloop setup PV + Tensor sV = + make_tensor(make_smem_ptr(storage.smem_v.data()), SmemLayoutV{}); + + Tensor tOsV = thr_mma_pv.partition_B(sV); // (MMA,MMA_N,MMA_K,PIPE) + Tensor tOrV = thr_mma_pv.make_fragment_B(tOsV); // (MMA,MMA_M,MMA_N,PIPE) + + auto problem_size_f = + apply_variable_length(problem_size, get<2, 0>(blk_coord)); + int k_tile_count = Fusion{}.get_unmasked_trip_count( + blk_coord, TileShape{}, problem_size_f); + + pipeline_q.consumer_wait(smem_pipe_read_q); + + // mapping into QK accumulator + Tensor cP = make_identity_tensor(take<0, 2>(TileShapeQK{})); + Tensor tPcP = thr_mma_qk.partition_C(cP); + int m_block = get<0>(wg_coord); + tPcP.data() = tPcP.data() + E<0>{} * m_block * get<0>(TileShapeQK{}); + + // Allocate PV acc + Tensor acc_pv = + partition_fragment_C(tiled_mma_pv, take<0, 2>(TileShapePV{})); + + cutlass::fmha::collective:: + CollectiveSoftmax + softmax{params}; + auto softmax_state = softmax.init(acc_pv, tiled_mma_pv); + + if (true) { + --k_tile_count; + // Allocate QK acc + Tensor acc_qk = + partition_fragment_C(tiled_mma_qk, take<0, 2>(TileShapeQK{})); + + pipeline.consumer_wait(smem_pipe_read); + math_wg_order_barrier.wait(); + + // MMA QK + warpgroup_fence_operand(acc_qk); + warpgroup_arrive(); + + gemm_zero_acc( + tiled_mma_qk, + tSrQ(_, _, _, smem_pipe_read_q.index()), + tSrK(_, _, _, smem_pipe_read.index()), + acc_qk); + warpgroup_commit_batch(); + math_wg_order_barrier.arrive(); + + ++smem_pipe_read; + + // Wait for the pipeline MMAs to drain + warpgroup_wait<0>(); + warpgroup_fence_operand(acc_qk); + + softmax.step(acc_qk, tiled_mma_qk, tPcP, softmax_state, problem_size_f); + + Tensor acc_qk_fixed = + make_acc_into_op(acc_qk, typename TiledMmaPV::LayoutA_TV{}); + + pipeline.consumer_wait(smem_pipe_read); + + // MMA PV + warpgroup_fence_operand(acc_pv); + warpgroup_fence_operand(acc_qk_fixed); + warpgroup_arrive(); + + gemm_zero_acc( + tiled_mma_pv, + acc_qk_fixed, + tOrV(_, _, _, smem_pipe_read.index()), + acc_pv); + warpgroup_commit_batch(); + + pipeline.consumer_release(smem_pipe_release); + ++smem_pipe_release; + + // Advance consumer pipeline + ++smem_pipe_read; + tPcP.data() = tPcP.data() + E<1>{} * get<1>(TileShapeQK{}); + } + + CUTLASS_PRAGMA_NO_UNROLL + while (k_tile_count > 0) { + --k_tile_count; + + // Allocate QK acc + Tensor acc_qk = + partition_fragment_C(tiled_mma_qk, take<0, 2>(TileShapeQK{})); + + pipeline.consumer_wait(smem_pipe_read); + + // MMA QK + warpgroup_fence_operand(acc_qk); + warpgroup_arrive(); + + gemm_zero_acc( + tiled_mma_qk, + tSrQ(_, _, _, smem_pipe_read_q.index()), + tSrK(_, _, _, smem_pipe_read.index()), + acc_qk); + warpgroup_commit_batch(); + + ++smem_pipe_read; + auto tok = pipeline.consumer_try_wait(smem_pipe_read); + + // Wait for the pipeline MMAs to drain + warpgroup_wait<0>(); + warpgroup_fence_operand(acc_qk); + warpgroup_fence_operand(acc_pv); + + if constexpr (kIsMainloopLocked) + math_wg_order_barrier.wait(); + softmax.template step( + acc_qk, + tiled_mma_qk, + tPcP, + softmax_state, + acc_pv, + tiled_mma_pv, + problem_size_f); + if constexpr (kIsMainloopLocked) + math_wg_order_barrier.arrive(); + + Tensor acc_qk_fixed = + make_acc_into_op(acc_qk, typename TiledMmaPV::LayoutA_TV{}); + + pipeline.consumer_wait(smem_pipe_read, tok); + + // MMA PV + warpgroup_fence_operand(acc_pv); + warpgroup_fence_operand(acc_qk_fixed); + warpgroup_arrive(); + + cute::gemm( + tiled_mma_pv, + acc_qk_fixed, + tOrV(_, _, _, smem_pipe_read.index()), + acc_pv); + warpgroup_commit_batch(); + + pipeline.consumer_release(smem_pipe_release); + ++smem_pipe_release; + + pipeline.consumer_release(smem_pipe_release); + ++smem_pipe_release; + + ++smem_pipe_read; + tPcP.data() = tPcP.data() + E<1>{} * get<1>(TileShapeQK{}); + } + + k_tile_count += + Fusion{}.get_masked_trip_count(blk_coord, TileShape{}, problem_size_f); + + CUTLASS_PRAGMA_NO_UNROLL + while (k_tile_count > 0) { + --k_tile_count; + + // Allocate QK acc + Tensor acc_qk = + partition_fragment_C(tiled_mma_qk, take<0, 2>(TileShapeQK{})); + + pipeline.consumer_wait(smem_pipe_read); + + // MMA QK + warpgroup_fence_operand(acc_qk); + warpgroup_arrive(); + + gemm_zero_acc( + tiled_mma_qk, + tSrQ(_, _, _, smem_pipe_read_q.index()), + tSrK(_, _, _, smem_pipe_read.index()), + acc_qk); + warpgroup_commit_batch(); + + ++smem_pipe_read; + auto tok = pipeline.consumer_try_wait(smem_pipe_read); + + // Wait for the pipeline MMAs to drain + warpgroup_wait<0>(); + warpgroup_fence_operand(acc_qk); + warpgroup_fence_operand(acc_pv); + + // if constexpr (kIsPersistent) + // if (k_tile_count == 0) + // pipeline_q.consumer_release(smem_pipe_release_q); + + if constexpr (kIsMainloopLocked) + math_wg_order_barrier.wait(); + softmax.step( + acc_qk, + tiled_mma_qk, + tPcP, + softmax_state, + acc_pv, + tiled_mma_pv, + problem_size_f); + if constexpr (kIsMainloopLocked) + math_wg_order_barrier.arrive(); + + Tensor acc_qk_fixed = + make_acc_into_op(acc_qk, typename TiledMmaPV::LayoutA_TV{}); + + pipeline.consumer_wait(smem_pipe_read, tok); + + // MMA PV + warpgroup_fence_operand(acc_pv); + warpgroup_fence_operand(acc_qk_fixed); + warpgroup_arrive(); + + cute::gemm( + tiled_mma_pv, + acc_qk_fixed, + tOrV(_, _, _, smem_pipe_read.index()), + acc_pv); + warpgroup_commit_batch(); + + pipeline.consumer_release(smem_pipe_release); + ++smem_pipe_release; + + pipeline.consumer_release(smem_pipe_release); + ++smem_pipe_release; + + ++smem_pipe_read; + tPcP.data() = tPcP.data() + E<1>{} * get<1>(TileShapeQK{}); + } + + if (kIsPersistent) + pipeline_q.consumer_release(smem_pipe_release_q); + + // Wait for the pipeline MMAs to drain + warpgroup_wait<0>(); + warpgroup_fence_operand(acc_pv); + + if (kIsPersistent) + pipeline.consumer_release(smem_pipe_release); + ++smem_pipe_release; + + Tensor lse = softmax.tail(softmax_state, acc_pv, tiled_mma_pv); + + return make_tuple(acc_pv, lse); + } +}; + +} // namespace cutlass::fmha::collective diff --git a/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_common.hpp b/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_common.hpp new file mode 100644 index 00000000..c0c2511b --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_common.hpp @@ -0,0 +1,317 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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/kernel_hardware_info.h" + +namespace cutlass::fmha::collective { + +using namespace cute; + +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; + if constexpr (rA == 2 && rB == 2 && rC == 1) { + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<1>(tA); k_block++) { + cute::gemm(atom, tA(_, k_block), tB(_, k_block), tC); + atom.accumulate_ = GMMA::ScaleOut::One; + } + } else { + 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_ = GMMA::ScaleOut::One; + } + } +} + +template +CUTE_DEVICE void gemm_zero_acc( + Atom& atom, + TA const& tA, + TB const& tB, + TC&& tC) { + atom.accumulate_ = GMMA::ScaleOut::Zero; + gemm_reset_zero_acc(atom, tA, tB, tC); +} + +template +CUTE_DEVICE constexpr typename T::value_type reduce(T const& t, Fn fn) { + if constexpr (decltype(size(t) % _2{} == _0{})::value) { + auto partial = make_tensor(size(t) / _2{}); + CUTE_UNROLL + for (int i = 0; i < size(partial); i++) { + partial(i) = fn(t(i), t(i + size(partial))); + } + return reduce(partial, fn); + } else { + auto result = t(_0{}); + CUTE_UNROLL + for (int i = 1; i < size(t); i++) { + result = fn(result, t(i)); + } + return result; + } +} + +struct fmha_max { + CUTE_DEVICE float operator()(float a, float b) { + return ::max(a, b); + } +}; + +template +inline auto __device__ constexpr layout_separate( + Threshold const& thr, + Source const& src, + Reference const& ref) { + auto lt = + filter(transform_layout(src, ref, [&](auto const& s, auto const& r) { + if constexpr (decltype(r < thr)::value) { + return s; + } else { + return make_layout(_1{}, _0{}); + } + })); + auto ge = + filter(transform_layout(src, ref, [&](auto const& s, auto const& r) { + if constexpr (decltype(r >= thr)::value) { + return s; + } else { + return make_layout(_1{}, _0{}); + } + })); + return make_layout(lt, ge); +} + +template +inline auto __device__ constexpr layout_acc_mn( + TiledMma const& tiled_mma, + Acc const& acc) { + auto separated = layout_separate( + get<0>(typename TiledMma::Shape_MNK{}), + get<0>(acc), + stride<1>(typename TiledMma::LayoutC_TV{})); + auto V_M = get<0>(separated); + auto V_N = get<1>(separated); + return make_layout( + make_layout(V_M, get<1>(acc)), make_layout(V_N, get<2>(acc))); +} + +template +inline auto __device__ constexpr layout_op_mk_v( + TiledMma const& tiled_mma, + Acc const& acc) { + return layout_separate( + get<0>(typename TiledMma::Shape_MNK{}), + get<0>(acc), + stride<1>(typename TiledMma::LayoutA_TV{})); +} + +template +inline auto __device__ constexpr tensor_op_mk_v( + TiledMma const& tiled_mma, + Acc&& acc) { + return make_tensor(acc.data(), layout_op_mk_v(tiled_mma, acc.layout())); +} + +template +inline auto __device__ constexpr reduction_target_n(TiledMma const& tiled_mma) { + auto separated = layout_separate( + get<0>(typename TiledMma::Shape_MNK{}), + make_layout(shape<0>(typename TiledMma::LayoutC_TV{})), + stride<0>(typename TiledMma::LayoutC_TV{})); + return get<1>(separated); +} + +template < + template < + cute::GMMA::Major, + cute::GMMA::Major, + cute::GMMA::ScaleIn, + cute::GMMA::ScaleIn> + class Primitive, + cute::GMMA::Major tA, + cute::GMMA::Major tB, + cute::GMMA::ScaleIn sA, + cute::GMMA::ScaleIn sB> +inline auto __device__ constexpr convert_to_gmma_rs( + cute::MMA_Atom> const& tiled_mma) { + using Atom = cute::MMA_Atom>; + using ElementA = typename Atom::ValTypeA; + using ElementB = typename Atom::ValTypeB; + using ElementC = typename Atom::ValTypeC; + using Shape_MNK = typename Atom::Shape_MNK; + using RS = decltype(cute::GMMA::rs_op_selector< + ElementA, + ElementB, + ElementC, + Shape_MNK, + tA, + tB, + sA, + sB>()); + return cute::MMA_Atom{}; +} + +template < + template + class Primitive, + cute::GMMA::ScaleIn sA, + cute::GMMA::ScaleIn sB> +inline auto __device__ constexpr convert_to_gmma_rs( + cute::MMA_Atom> const& tiled_mma) { + using Atom = cute::MMA_Atom>; + using ElementA = typename Atom::ValTypeA; + using ElementB = typename Atom::ValTypeB; + using ElementC = typename Atom::ValTypeC; + using Shape_MNK = typename Atom::Shape_MNK; + constexpr auto tA = cute::GMMA::Major::K; + constexpr auto tB = cute::GMMA::Major::K; + using RS = decltype(cute::GMMA::rs_op_selector< + ElementA, + ElementB, + ElementC, + Shape_MNK, + tA, + tB, + sA, + sB>()); + return cute::MMA_Atom{}; +} + +template +CUTE_DEVICE auto constexpr convert_to_gmma_rs( + cute::TiledMMA const& tiled_mma) { + return cute::TiledMMA{}; +} + +template +CUTE_DEVICE auto constexpr convert_c_layout_to_a_layout( + CLayout const& c, + AValueShape const& a) { + return make_layout( + make_shape(a, shape<1>(c), make_shape(shape<2>(c), size<0>(c) / size(a))), + make_stride( + stride<0>(c), + stride<1>(c), + make_stride(stride<2>(c), size<2>(a) * stride<0, 2>(c)))); +} + +template +CUTE_DEVICE constexpr auto unstageSmemLayout( + Layout const& layout, + Stages stages = {}) { + return composition(layout, make_tuple(_, _, make_layout(stages))); +} + +template +CUTE_DEVICE auto make_acc_into_op( + Accumulator const& acc, + OperandLayout_TV const& operand_layout_tv) { + Tensor operand = make_fragment_like( + convert_c_layout_to_a_layout(acc.layout(), shape<1>(operand_layout_tv))); + Tensor operand_as_acc = make_tensor(operand.data(), acc.layout()); + + cute::copy(acc, operand_as_acc); + + if constexpr (sizeof(Element) == 1) { + // 00 11 22 33 00 11 22 33 acc layout + // 00 00 11 11 22 22 33 33 operand layout + // BB AA AA BB AA BB BB AA conflict-free exchange pattern + // 16-bit exchange; so process two at a time + // potentially + int tid = threadIdx.x % 4; + auto values_u32 = recast(operand); + + CUTE_UNROLL + for (int n = 0; n < size<1>(values_u32); n++) { + CUTE_UNROLL + for (int k = 0; k < size<2>(values_u32); k++) { + CUTE_UNROLL + for (int ii = 0; ii < 8; ii += 4) { + uint32_t values_tmp_0 = values_u32(ii / 2 + 0, n, k); + uint32_t values_tmp_1 = values_u32(ii / 2 + 1, n, k); + + // step A: + // t 1 v 0 -> t 0 v 1 + // t 2 v 0 -> t 1 v 0 + // t 0 v 1 -> t 2 v 0 + // t 3 v 1 -> t 3 v 1 + + int v_to_send = tid == 1 || tid == 2 ? 0 : 1; + int v_to_recv = v_to_send; + int t_to_recv_from = (0x3021 >> (tid * 4)) & 0xF; + + uint32_t values_tmp_a = v_to_send == 0 ? values_tmp_0 : values_tmp_1; + + values_tmp_a = + __shfl_sync(0xFFFFFFFF, values_tmp_a, t_to_recv_from, 4); + + // step B: + // t 0 v 0 -> t 0 v 0 + // t 3 v 0 -> t 1 v 1 + // t 1 v 1 -> t 2 v 1 + // t 2 v 1 -> t 3 v 0 + + v_to_send = 1 - v_to_send; + v_to_recv = 1 - v_to_recv; + t_to_recv_from = (0x2130 >> (tid * 4)) & 0xF; + + uint32_t values_tmp_b = v_to_send == 0 ? values_tmp_0 : values_tmp_1; + + values_tmp_b = + __shfl_sync(0xFFFFFFFF, values_tmp_b, t_to_recv_from, 4); + + values_u32(ii / 2 + 0, n, k) = __byte_perm( + values_tmp_a, values_tmp_b, v_to_send == 0 ? 0x1054 : 0x5410); + values_u32(ii / 2 + 1, n, k) = __byte_perm( + values_tmp_a, values_tmp_b, v_to_send == 0 ? 0x3276 : 0x7632); + } + } + } + } + + return operand; +} + +} // namespace cutlass::fmha::collective diff --git a/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_epilogue.hpp b/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_epilogue.hpp new file mode 100644 index 00000000..1c18e861 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_epilogue.hpp @@ -0,0 +1,239 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/collective/default_epilogue.hpp" +#include "cutlass/epilogue/thread/linear_combination.h" + +#include "natten/cuda/fmha_hopper/collective/fmha_common.hpp" +#include "natten/cuda/fmha_hopper/collective/fmha_varlen.hpp" + +namespace cutlass::fmha::collective { + +template < + class Element, + class ElementAccumulator, + class TileShape_WG, + bool IsVarlen> +struct FmhaFwdEpilogueSm90 { + static constexpr int Alignment = 16 / sizeof(Element); + + using DefaultOperation = cutlass::epilogue::fusion:: + LinearCombination; + + // Varlen can't use TMA store in SM90. Needs to fall back to STG. + using EpilogueType = std::conditional_t< + IsVarlen, + cutlass::epilogue::NoSmemWarpSpecialized, + cutlass::epilogue::TmaWarpSpecialized>; + + using CollectiveEpilogueTMA = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm90, + cutlass::arch::OpClassTensorOp, + TileShape_WG, + Shape<_1, _1, _1>, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, + ElementAccumulator, + void, + cute::tuple>, + Alignment, + Element, + cute::tuple>, + Alignment, + EpilogueType, + DefaultOperation>::CollectiveOp; + + struct Arguments { + Element* ptr_O; + cute::tuple> dO; + + ElementAccumulator* ptr_LSE; + cute::tuple> dLSE; + }; + + struct Params { + ElementAccumulator* ptr_LSE; + cute::tuple> dLSE; + + typename CollectiveEpilogueTMA::Params epilogue_TMA; + }; + + using TensorStorage = typename CollectiveEpilogueTMA::TensorStorage; + using PipelineStorage = typename CollectiveEpilogueTMA::PipelineStorage; + using LoadPipeline = typename CollectiveEpilogueTMA::LoadPipeline; + static constexpr int TmaTransactionBytes = + CollectiveEpilogueTMA::TmaTransactionBytes; + + template + static Params to_underlying_arguments( + ProblemShape const& problem_size, + Arguments const& args, + void* workspace = nullptr) { + auto problem_size_o = apply_variable_length_scheduler(make_shape( + get<2>(problem_size), + get<4>(problem_size), + 1, + make_shape(get<0>(problem_size), get<1>(problem_size)))); + typename CollectiveEpilogueTMA::Arguments args_tma{ + {}, args.ptr_O, args.dO, args.ptr_O, args.dO}; + return Params{ + args.ptr_LSE, + args.dLSE, + CollectiveEpilogueTMA::to_underlying_arguments( + problem_size_o, args_tma, workspace)}; + } + + template < + class TileShape, + class BlkCoord, + class ResultTuple, + class TiledMma, + class ProblemShape> + CUTLASS_DEVICE void operator()( + TileShape const& tile_shape, + BlkCoord const& blk_coord_, + ResultTuple const& result, + TiledMma const& tiled_mma, + ProblemShape const& problem_size, + Params const& params, + LoadPipeline epi_load_pipeline, + TensorStorage& epi_tensor_storage) { + using X = Underscore; + + auto acc = get<0>(result); + auto lse = get<1>(result); + + auto thr_mma = tiled_mma.get_thread_slice(threadIdx.x); + + int seq_off = 0; + int seqlen_q = get<2>(problem_size); + int num_batch = get<0>(problem_size); + + BlkCoord blk_coord = blk_coord_; + if constexpr (is_problem_shape_variable_length(ProblemShape{})) { + auto cumulative_length = get<2>(problem_size).cumulative_length; + if (cumulative_length != nullptr) { + seq_off = cumulative_length[get<2, 0>(blk_coord)]; + seqlen_q = cumulative_length[get<2, 0>(blk_coord) + 1] - seq_off; + num_batch = 1; + get<2, 0>(blk_coord) = 0; + } + } + + int num_heads = get<1>(problem_size); + // Epilogue for lse + Tensor mLSE_ = make_tensor( + make_gmem_ptr(params.ptr_LSE), + make_shape( + seqlen_q, get<1>(tile_shape), make_shape(num_batch, num_heads)), + make_stride( + get<0>(params.dLSE), + _0{}, + make_stride(get<1, 0>(params.dLSE), _1{}))); + Tensor mLSE = + domain_offset(make_coord(seq_off, _0{}, make_coord(_0{}, _0{})), mLSE_); + Tensor gLSE_full = + local_tile(mLSE, tile_shape, make_coord(_, _, _), Step<_1, _1, X>{}); + Tensor gLSE = gLSE_full( + _, _, get<0>(blk_coord), get<1>(blk_coord), get<2>(blk_coord)); + Tensor tOgLSE = thr_mma.partition_C(gLSE); + Tensor cO = make_identity_tensor(take<0, 2>(tile_shape)); + Tensor tOcO = thr_mma.partition_C(cO); + if (get<1>(tOcO(_0{})) == 0) { + auto tOgLSE_mn = + make_tensor(tOgLSE.data(), layout_acc_mn(tiled_mma, tOgLSE.layout())); + auto tOcO_mn = + make_tensor(tOcO.data(), layout_acc_mn(tiled_mma, tOcO.layout())); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<0>(tOgLSE_mn); i++) { + if (get<0>(tOcO_mn(i)) + get<0>(blk_coord) * get<0>(tile_shape) < + seqlen_q) { + tOgLSE_mn(i, _0{}) = lse(i); + } + } + } + auto problem_size_o = make_shape( + seqlen_q, + get<4>(problem_size), + _, + make_shape(num_batch, get<1>(problem_size))); + + auto epilogue = [&]() { + if constexpr (IsVarlen) { + auto epilogue_params = params.epilogue_TMA; + epilogue_params.ptr_D += get<0>(epilogue_params.dD) * seq_off; + return CollectiveEpilogueTMA(epilogue_params, epi_tensor_storage); + } else { + return CollectiveEpilogueTMA(params.epilogue_TMA, epi_tensor_storage); + } + }(); + + using EpiStorePipeline = typename CollectiveEpilogueTMA::StorePipeline; + typename EpiStorePipeline::Params epi_store_pipeline_params; + epi_store_pipeline_params.always_wait = true; + EpiStorePipeline epi_store_pipeline(epi_store_pipeline_params); + + typename CollectiveEpilogueTMA::LoadPipelineState + epi_load_pipe_consumer_state; + PipelineState epi_store_pipe_producer_state = + cutlass::make_producer_start_state(); + + auto + [epi_load_pipe_consumer_state_next, + epi_store_pipe_producer_state_next] = + epilogue.store( + epi_load_pipeline, + epi_load_pipe_consumer_state, + epi_store_pipeline, + epi_store_pipe_producer_state, + problem_size_o, + tile_shape, + make_coord(get<0>(blk_coord), _0{}, _, get<2>(blk_coord)), + acc, + tiled_mma, + threadIdx.x % cutlass::NumThreadsPerWarpGroup, + epi_tensor_storage); + + epilogue.store_tail( + epi_load_pipeline, + epi_load_pipe_consumer_state_next, + epi_store_pipeline, + epi_store_pipe_producer_state_next); + } +}; + +} // namespace cutlass::fmha::collective diff --git a/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_epilogue_bwd.hpp b/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_epilogue_bwd.hpp new file mode 100644 index 00000000..085aa6fc --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_epilogue_bwd.hpp @@ -0,0 +1,259 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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/epilogue/collective/default_epilogue.hpp" + +#include "natten/cuda/fmha_hopper/collective/fmha_epilogue.hpp" +#include "natten/cuda/fmha_hopper/collective/fmha_varlen.hpp" + +namespace cutlass::fmha::collective { + +template < + class Element, + class ElementAccumulator, + class TileShape_WG, + bool IsVarlen> +struct FmhaBwdEpilogueKV { + static constexpr int Alignment = 16 / sizeof(Element); + + struct Arguments { + Element* ptr_K; + cute::tuple dK; + + Element* ptr_V; + cute::tuple dV; + }; + + // Varlen can't use TMA store in SM90. Needs to fall back to STG. + using EpilogueType = std::conditional_t< + IsVarlen, + cutlass::epilogue::NoSmemWarpSpecialized, + cutlass::epilogue::TmaWarpSpecialized>; + + // NOTE (ahassani): standard LinearCombination (axpby) was commented out in + // favor of the EVT fusion, but as far as I can tell they're doing the same + // thing, except the latter will fail dispatching to nosmem epilogue + // (non-TMA). In case there is a reason why this was done, I'm making the + // fusion also conditional on varlen. + using DefaultOperationNonTMA = cutlass::epilogue::fusion:: + LinearCombination; + + using DefaultOperationTMA = cutlass::epilogue::fusion::Sm90EVT< + cutlass::epilogue::fusion::Sm90Compute< + cutlass::first, + Element, + ElementAccumulator, + cutlass::FloatRoundStyle::round_to_nearest>, + cutlass::epilogue::fusion::Sm90AccFetch>; + + using DefaultOperation = + std::conditional_t; + + using CollectiveEpilogueTMA = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm90, + cutlass::arch::OpClassTensorOp, + TileShape_WG, + Shape<_1, _1, _1>, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, + ElementAccumulator, + void, + cute::tuple>, + Alignment, + Element, + cute::tuple>, + Alignment, + EpilogueType, + DefaultOperation>::CollectiveOp; + + struct Params { + typename CollectiveEpilogueTMA::Params epilogue_K; + typename CollectiveEpilogueTMA::Params epilogue_V; + }; + + using TensorStorage = typename CollectiveEpilogueTMA::TensorStorage[2]; + using PipelineStorage = typename CollectiveEpilogueTMA::PipelineStorage; + using LoadPipeline = typename CollectiveEpilogueTMA::LoadPipeline; + static constexpr int TmaTransactionBytes = + CollectiveEpilogueTMA::TmaTransactionBytes; + + template + static Params to_underlying_arguments( + ProblemShape const& problem_size, + Arguments const& args, + void* workspace = nullptr) { + auto dK = make_stride( + get<2>(args.dK), + get<3>(args.dK), + make_stride(get<0>(args.dK), get<1>(args.dK))); + auto dV = make_stride( + get<2>(args.dV), + get<3>(args.dV), + make_stride(get<0>(args.dV), get<1>(args.dV))); + + auto problem_size_kv = apply_variable_length_scheduler(make_shape( + get<3>(problem_size), + get<4>(problem_size), + 1, + make_shape(get<0>(problem_size), get<1>(problem_size)))); + typename CollectiveEpilogueTMA::Arguments args_k{ + {}, args.ptr_K, dK, args.ptr_K, dK}; + typename CollectiveEpilogueTMA::Arguments args_v{ + {}, args.ptr_V, dV, args.ptr_V, dV}; + return Params{ + CollectiveEpilogueTMA::to_underlying_arguments( + problem_size_kv, args_k, nullptr), + CollectiveEpilogueTMA::to_underlying_arguments( + problem_size_kv, args_v, nullptr)}; + } + + template < + class TileShape, + class BlkCoord, + class ResultTuple, + class TiledMma, + class ProblemShape> + CUTLASS_DEVICE void operator()( + TileShape const& tile_shape, + BlkCoord const& blk_coord_, + ResultTuple const& result, + TiledMma const& tiled_mma, + ProblemShape const& problem_size, + Params const& params, + LoadPipeline epi_load_pipeline, + TensorStorage& epi_tensor_storage) { + auto acc_k = get<0>(result); + auto acc_v = get<1>(result); + + int seq_off = 0; + int seqlen = get<3>(problem_size); + int num_batch = get<0>(problem_size); + + BlkCoord blk_coord = blk_coord_; + if constexpr (is_problem_shape_variable_length(ProblemShape{})) { + auto cumulative_length = get<3>(problem_size).cumulative_length; + if (cumulative_length != nullptr) { + seq_off = cumulative_length[get<2, 0>(blk_coord)]; + seqlen = cumulative_length[get<2, 0>(blk_coord) + 1] - seq_off; + num_batch = 1; + get<2, 0>(blk_coord) = 0; + } + } + + auto problem_size_kv = make_shape( + seqlen, + get<4>(problem_size), + _, + make_shape(num_batch, get<1>(problem_size))); + + using EpiStorePipeline = typename CollectiveEpilogueTMA::StorePipeline; + typename EpiStorePipeline::Params epi_store_pipeline_params; + epi_store_pipeline_params.always_wait = true; + EpiStorePipeline epi_store_pipeline(epi_store_pipeline_params); + + typename CollectiveEpilogueTMA::LoadPipelineState + epi_load_pipe_consumer_state; + PipelineState epi_store_pipe_producer_state = + cutlass::make_producer_start_state(); + + auto [epilogue_k, epilogue_v] = [&]() { + if constexpr (IsVarlen) { + auto params_K = params.epilogue_K; + auto params_V = params.epilogue_V; + + params_K.ptr_D += get<0>(params_K.dD) * seq_off; + params_V.ptr_D += get<0>(params_V.dD) * seq_off; + return make_tuple( + CollectiveEpilogueTMA(params_K, epi_tensor_storage[0]), + CollectiveEpilogueTMA(params_V, epi_tensor_storage[1])); + } else { + return make_tuple( + CollectiveEpilogueTMA(params.epilogue_K, epi_tensor_storage[0]), + CollectiveEpilogueTMA(params.epilogue_V, epi_tensor_storage[1])); + } + }(); + + { + auto + [epi_load_pipe_consumer_state_next, + epi_store_pipe_producer_state_next] = + epilogue_k.store( + epi_load_pipeline, + epi_load_pipe_consumer_state, + epi_store_pipeline, + epi_store_pipe_producer_state, + problem_size_kv, + tile_shape, + make_coord(get<1>(blk_coord), _0{}, _, get<2>(blk_coord)), + acc_k, + tiled_mma, + threadIdx.x % cutlass::NumThreadsPerWarpGroup, + epi_tensor_storage[0]); + } + + { + auto + [epi_load_pipe_consumer_state_next, + epi_store_pipe_producer_state_next] = + epilogue_v.store( + epi_load_pipeline, + epi_load_pipe_consumer_state, + epi_store_pipeline, + epi_store_pipe_producer_state, + problem_size_kv, + tile_shape, + make_coord(get<1>(blk_coord), _0{}, _, get<2>(blk_coord)), + acc_v, + tiled_mma, + threadIdx.x % cutlass::NumThreadsPerWarpGroup, + epi_tensor_storage[1]); + + epilogue_k.store_tail( + epi_load_pipeline, + epi_load_pipe_consumer_state_next, + epi_store_pipeline, + epi_store_pipe_producer_state_next); + + epilogue_v.store_tail( + epi_load_pipeline, + epi_load_pipe_consumer_state_next, + epi_store_pipeline, + epi_store_pipe_producer_state_next); + } + } +}; + +} // namespace cutlass::fmha::collective diff --git a/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_fusion.hpp b/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_fusion.hpp new file mode 100644 index 00000000..59e84684 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_fusion.hpp @@ -0,0 +1,281 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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 DefaultFusion { + template + CUTLASS_DEVICE int get_trip_count( + BlkCoord const& blk_coord, + TileShape const& tile_shape, + ProblemSize const& problem_size) { + return ceil_div(get<3>(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 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 0; + } + + template + CUTLASS_DEVICE void before_softmax( + AccQK& acc_qk, + IndexQK const& index_qk, + ProblemSize const& problem_size + + ) { + return; + } +}; + +struct ResidualFusion : DefaultFusion { + using Base = DefaultFusion; + + template + CUTLASS_DEVICE int get_masked_trip_count( + BlkCoord const& blk_coord, + TileShape const& tile_shape, + ProblemSize const& problem_size) { + return 1; + } + + 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) - 1; + } + + template + CUTLASS_DEVICE void before_softmax( + 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<3>(problem_size)) { + acc_qk(i) = -INFINITY; + } + } + } +}; + +struct CausalFusion : DefaultFusion { + using Base = DefaultFusion; + + 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 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), 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 ceil_div(get<0>(tile_shape), get<1>(tile_shape)); + } + + 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 before_softmax( + 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) + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(acc_qk); i++) { + auto pos = index_qk(i); + if (get<0>(pos) < get<1>(pos) || get<1>(pos) >= get<3>(problem_size)) { + acc_qk(i) = -INFINITY; + } + } + } +}; + +template +struct FusionBwdAdapter { + template + CUTLASS_DEVICE int get_trip_count( + BlkCoord const& blk_coord, + TileShape const& tile_shape, + ProblemSize const& problem_size) { + return Base{}.get_trip_count( + select<1, 0, 2>(blk_coord), + select<1, 0, 2>(tile_shape), + select<0, 1, 3, 2, 4>(problem_size)); + } + + template + CUTLASS_DEVICE void before_softmax( + AccQK& acc_qk, + IndexQK const& index_qk, + ProblemSize const& problem_size) { + if constexpr (IsVarlen || is_same_v) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(acc_qk); i++) { + auto pos = index_qk(i); + if (get<0>(pos) < 0 || get<1>(pos) < 0 || + get<0>(pos) >= get<3>(problem_size) || + get<1>(pos) >= get<2>(problem_size)) { + acc_qk(i) = -INFINITY; + } + } + } else if constexpr (is_same_v) { + // no op + } else { + // Fallthrough in case we add more fmha masks + auto index_base = index_qk(_0{}); + auto index_shape = shape(index_qk); + auto index_stride = transform_leaf(stride(index_qk), [](auto elem) { + if constexpr (is_scaled_basis::value) { + if constexpr (decltype(elem.mode() == _0{})::value) { + return ScaledBasis(elem.value()); + } else { + return ScaledBasis(elem.value()); + } + } else { + return elem; + } + }); + auto index_qk_bwd = make_tensor( + make_inttuple_iter(select<1, 0>(index_base)), + make_layout(index_shape, index_stride)); + Base{}.before_softmax(acc_qk, index_qk_bwd, problem_size); + } + } + + template + CUTLASS_DEVICE bool is_contributing( + BlkCoord const& blk_coord, + TileShape const& tile_shape, + ProblemSize const& problem_size) { + return true; + } +}; + +template <> +struct FusionBwdAdapter { + template + CUTLASS_DEVICE int get_trip_count( + BlkCoord const& blk_coord, + TileShape const& tile_shape, + ProblemSize const& problem_size) { + return ceil_div(get<2>(problem_size), get<0>(TileShape{})); + } + + template + CUTLASS_DEVICE void before_softmax( + AccQK& acc_qk, + IndexQK const& index_qk, + ProblemSize const& problem_size + + ) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(acc_qk); i++) { + auto pos = index_qk(i); + // Varlen needs additional boundary checks + if constexpr (IsVarlen) { + if (get<1>(pos) < get<0>(pos) || get<1>(pos) >= get<2>(problem_size) || + 0 > get<0>(pos) || get<0>(pos) >= get<3>(problem_size) || + 0 > get<1>(pos)) { + acc_qk(i) = -INFINITY; + } + } else { + if (get<1>(pos) < get<0>(pos) || get<1>(pos) >= get<2>(problem_size)) { + acc_qk(i) = -INFINITY; + } + } + } + } + + template + CUTLASS_DEVICE bool is_contributing( + BlkCoord const& blk_coord, + TileShape const& tile_shape, + ProblemSize const& problem_size) { + int max_q = get<0>(blk_coord) * get<0>(tile_shape) + get<0>(tile_shape); + int min_k = get<1>(blk_coord) * get<1>(tile_shape); + return min_k <= max_q; + } +}; + +} // namespace cutlass::fmha::collective diff --git a/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_varlen.hpp b/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_varlen.hpp new file mode 100644 index 00000000..b70bffd6 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_hopper/collective/fmha_varlen.hpp @@ -0,0 +1,109 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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 VariableLength { + int max_length; + int* cumulative_length = nullptr; + int total_length = -1; + + CUTE_HOST_DEVICE operator int() const { + return max_length; + } +}; + +template +struct is_variable_length_impl : std::false_type {}; +template <> +struct is_variable_length_impl : std::true_type {}; +template +constexpr bool is_variable_length_v = + is_variable_length_impl>::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.cumulative_length[idx + 1] - s.cumulative_length[idx]; + } else { + return s; + } + }); +} + +template +CUTE_HOST_DEVICE constexpr auto apply_variable_length_scheduler( + Shape const& shape) { + return transform_leaf(shape, [&](auto const& s) { + if constexpr (is_variable_length_v) { + return s.total_length; + } else { + return s; + } + }); +} + +template +CUTE_HOST_DEVICE constexpr bool is_problem_shape_variable_length( + Shape const& shape) { + bool is_varlen = false; + cute::for_each(shape, [&](auto const& s) { + if constexpr (is_variable_length_v) { + is_varlen = true; + } + }); + return is_varlen; +} + +} // namespace cutlass::fmha::collective + +namespace cute { + +template <> +struct is_integral : true_type {}; + +CUTE_HOST_DEVICE +void print(cutlass::fmha::collective::VariableLength a) { + printf("Varlen<%d, %p>", a.max_length, a.cumulative_length); +} + +} // namespace cute diff --git a/natten/csrc/include/natten/cuda/fmha_hopper/device/fmha_bwd_sm90.hpp b/natten/csrc/include/natten/cuda/fmha_hopper/device/fmha_bwd_sm90.hpp new file mode 100644 index 00000000..ea544023 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_hopper/device/fmha_bwd_sm90.hpp @@ -0,0 +1,400 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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 + +/*! + \file + \brief An universal device layer for cutlass 3.x-style kernels. +*/ + +#include +// common +#include "cutlass/cutlass.h" + +#include "natten/cuda/fmha_hopper/collective/fmha_collective_bwd_tma_warpspecialized.hpp" +#include "natten/cuda/fmha_hopper/collective/fmha_epilogue_bwd.hpp" +#include "natten/cuda/fmha_hopper/collective/fmha_fusion.hpp" +#include "natten/cuda/fmha_hopper/collective/fmha_varlen.hpp" +#include "natten/cuda/fmha_hopper/device/fmha_sm90.hpp" +#include "natten/cuda/fmha_hopper/kernel/fmha_kernel_bwd_convert.hpp" +#include "natten/cuda/fmha_hopper/kernel/fmha_kernel_bwd_sum_OdO.hpp" +#include "natten/cuda/fmha_hopper/kernel/fmha_kernel_tma_warpspecialized.hpp" +#include "natten/cuda/fmha_hopper/kernel/fmha_tile_scheduler.hpp" + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::fmha::device { + +//////////////////////////////////////////////////////////////////////////////// +////////////////////////////// CUTLASS 3.x API ///////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +template < + class ProblemShape, + class Element, + class ElementAccumulator, + class TileShape, + class Fusion, + class... Options> +class FmhaBwdSm90 { + public: + /// Argument structure: User API + struct Arguments { + ProblemShape problem_size; + + const Element* ptr_Q; + cute::tuple stride_Q; + const Element* ptr_K; + cute::tuple stride_K; + const Element* ptr_V; + cute::tuple stride_V; + + const Element* ptr_O; + cute::tuple stride_O; + const ElementAccumulator* ptr_LSE; + // NATTEN has a different LSE layout + cute::tuple stride_LSE; + + const Element* ptr_dO; + cute::tuple stride_dO; + + Element* ptr_dQ; + cute::tuple stride_dQ; + Element* ptr_dK; + cute::tuple stride_dK; + Element* ptr_dV; + cute::tuple stride_dV; + + // if zero, defaults to 1/sqrt(D) + float scale_softmax = 0.0f; + + cutlass::KernelHardwareInfo hw_info; + }; + + using OperationSumOdO = cutlass::fmha::device::Sm90DeviceKernel< + cutlass::fmha::kernel:: + FmhaKernelBwdSumOdO>; + using OperationConvert = cutlass::fmha::device::Sm90DeviceKernel< + cutlass::fmha::kernel:: + FmhaKernelBwdConvert>; + + using Mainloop = + cutlass::fmha::collective::FmhaBwdMainloopTmaWarpSpecializedSm90< + Element, + ElementAccumulator, + TileShape, + cutlass::fmha::collective::FusionBwdAdapter, + Options...>; + + static constexpr bool IsVarlen = + cutlass::fmha::collective::is_problem_shape_variable_length( + ProblemShape{}); + using Epilogue = cutlass::fmha::collective::FmhaBwdEpilogueKV< + Element, + ElementAccumulator, + typename Mainloop::TileShapePV, + IsVarlen>; + + using Operation = cutlass::fmha::device::Sm90DeviceKernel< + cutlass::fmha::kernel::FmhaKernelTmaWarpSpecialized< + ProblemShape, + Mainloop, + Epilogue, + cutlass::fmha::kernel::TileSchedulerBwdAdapter< + cutlass::fmha::kernel::IndividualTileScheduler>, + Options...>>; + + struct Params { + OperationSumOdO op_sum_OdO; + Operation op; + OperationConvert op_convert; + ElementAccumulator* dQ_acc; + size_t dQ_acc_size; + }; + + private: + Params params_; + + // Scaled LSE and sum OdO are batch packed, not sequence packed, to avoid + // alignment issues with LSE. + static typename OperationSumOdO::Arguments to_sum_OdO_arguments( + Arguments const& args, + ElementAccumulator* sum_OdO = nullptr, + ElementAccumulator* scaled_lse = nullptr) { + using namespace cute; + auto [B, H, Q_, K, D] = args.problem_size; + D = cutlass::round_up(D, 8); // Alignment + int Q = cutlass::round_up(static_cast(Q_), 8); // Alignment + auto stride_scalar = + make_stride(static_cast(H) * static_cast(Q), Q, _1{}); + // auto log2_e = log2f(expf(1.0f)); + return typename OperationSumOdO::Arguments{ + args.problem_size, + args.ptr_O, + args.stride_O, + args.ptr_dO, + args.stride_dO, + sum_OdO, + stride_scalar, + args.ptr_LSE, + args.stride_LSE, + scaled_lse, + stride_scalar/*, + 1.0f, + log2_e*/}; + } + + // F32 dQ is batch-packed, not sequence-packed + static typename OperationConvert::Arguments to_convert_arguments( + Arguments const& args, + ElementAccumulator* src = nullptr) { + auto [B, H, Q_, K, D] = args.problem_size; + D = cutlass::round_up(D, 8); // Alignment + int Q = cutlass::round_up(static_cast(Q_), 8); // Alignment + auto stride_src_dQ = make_stride( + static_cast(H * D) * static_cast(Q), Q * D, D, _1{}); + return typename OperationConvert::Arguments{ + args.problem_size, + src, + stride_src_dQ, + nullptr, + stride_src_dQ, + nullptr, + stride_src_dQ, + args.ptr_dQ, + args.stride_dQ, + nullptr, + args.stride_dK, + nullptr, + args.stride_dV}; + } + + static typename Operation::Arguments to_bwd_arguments( + Arguments const& args, + ElementAccumulator* sum_OdO = nullptr, + cute::tuple const& stride_sum_OdO = {}, + ElementAccumulator* scaled_lse = nullptr, + cute::tuple const& stride_scaled_lse = {}, + ElementAccumulator* dQ_acc = nullptr, + cute::tuple const& stride_dQ = {}) { + return typename Operation::Arguments{ + args.problem_size, + {args.ptr_Q, + args.stride_Q, + args.ptr_K, + args.stride_K, + args.ptr_V, + args.stride_V, + args.ptr_dO, + args.stride_dO, + scaled_lse, + stride_scaled_lse, + sum_OdO, + stride_sum_OdO, + dQ_acc, + stride_dQ, + args.scale_softmax}, + {args.ptr_dK, args.stride_dK, args.ptr_dV, args.stride_dV}, + args.hw_info}; + } + + public: + /// Determines whether the GEMM can execute the given problem. + static Status can_implement(Arguments const& args) { + Status status = Status::kSuccess; + + status = OperationSumOdO::can_implement(to_sum_OdO_arguments(args)); + if (status != Status::kSuccess) { + return status; + } + + status = OperationConvert::can_implement(to_convert_arguments(args)); + if (status != Status::kSuccess) { + return status; + } + + status = Operation::can_implement(to_bwd_arguments(args)); + if (status != Status::kSuccess) { + return status; + } + + return status; + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const& args) { + auto [B_, H_, Q_, K, D_] = args.problem_size; + size_t B = static_cast(B_); + size_t H = static_cast(H_); + size_t D = cutlass::round_up(static_cast(D_), 8); // Alignment + size_t Q = cutlass::round_up(static_cast(Q_), 8); // Alignment + size_t workspace_bytes = 0; + // All three intermediary tensors are batch-packed, not sequence packed + // OdO vector + workspace_bytes += B * H * Q * sizeof(ElementAccumulator); + // scaled LSE vector + workspace_bytes += B * H * Q * sizeof(ElementAccumulator); + // FP32 versions of outputs that are churned (start off with Q only) + workspace_bytes += B * H * Q * D * sizeof(ElementAccumulator); + return workspace_bytes; + } + + /// Initializes state from arguments. + Status initialize_split( + Arguments const& args, + void* workspace_dQ, + void* workspace_sum_OdO, + void* workspace_scaled_lse, + cudaStream_t stream = nullptr) { + CUTLASS_TRACE_HOST( + "FmhaBwdSm90::initialize_split() - workspace_dQ=" + << workspace_dQ << ", workspace_sum_OdO=" << workspace_sum_OdO + << "stream: " << (stream ? "non-null" : "null")); + + auto [B_, H_, Q_, K, D_] = args.problem_size; + size_t B = static_cast(B_); + size_t H = static_cast(H_); + size_t D = cutlass::round_up(static_cast(D_), 8); // Alignment + size_t Q = cutlass::round_up(static_cast(Q_), 8); // Alignment + ElementAccumulator* sum_OdO = + reinterpret_cast(workspace_sum_OdO); + ElementAccumulator* scaled_lse = + reinterpret_cast(workspace_scaled_lse); + ElementAccumulator* dQ_acc = + reinterpret_cast(workspace_dQ); + params_.dQ_acc = dQ_acc; + params_.dQ_acc_size = B * H * Q * D * sizeof(ElementAccumulator); + auto args_sum_OdO = to_sum_OdO_arguments(args, sum_OdO, scaled_lse); + auto args_convert = to_convert_arguments(args, dQ_acc); + params_.op_sum_OdO.initialize(args_sum_OdO, nullptr, stream); + params_.op_convert.initialize(args_convert, nullptr, stream); + auto args_bwd = to_bwd_arguments( + args, + sum_OdO, + args_sum_OdO.stride_sum_OdO, + scaled_lse, + args_sum_OdO.stride_scaled_lse, + dQ_acc, + args_convert.stride_src_dQ); + params_.op.initialize(args_bwd, nullptr, stream); + + return Status::kSuccess; + } + + /// Initializes state from arguments. + Status initialize( + Arguments const& args, + void* workspace = nullptr, + cudaStream_t stream = nullptr) { + CUTLASS_TRACE_HOST( + "FmhaBwdSm90::initialize() - workspace " + << workspace << ", stream: " << (stream ? "non-null" : "null")); + + auto [B_, H_, Q_, K, D_] = args.problem_size; + size_t B = static_cast(B_); + size_t H = static_cast(H_); + size_t D = cutlass::round_up(static_cast(D_), 8); // Alignment + size_t Q = cutlass::round_up(static_cast(Q_), 8); // Alignment + char* workspace_chr = reinterpret_cast(workspace); + ElementAccumulator* sum_OdO = + reinterpret_cast(workspace_chr); + workspace_chr += B * H * Q * sizeof(ElementAccumulator); + ElementAccumulator* scaled_lse = + reinterpret_cast(workspace_chr); + workspace_chr += B * H * Q * sizeof(ElementAccumulator); + ElementAccumulator* dQ_acc = + reinterpret_cast(workspace_chr); + return initialize_split(args, dQ_acc, sum_OdO, scaled_lse, stream); + } + + /// 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("FmhaBwdSm90::run()"); + + Status result = Status::kSuccess; + result = params.op_sum_OdO.run(stream); + if (result != Status::kSuccess) { + return result; + } + + auto cuda_result = + cudaMemsetAsync(params.dQ_acc, 0, params.dQ_acc_size, stream); + if (cuda_result != cudaSuccess) { + return Status::kErrorInternal; + } + + result = params.op.run(stream); + if (result != Status::kSuccess) { + return result; + } + + result = params.op_convert.run(stream); + if (result != Status::kSuccess) { + return result; + } + + 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; + } + + /// 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); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::fmha::device + +//////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fmha_hopper/device/fmha_sm90.hpp b/natten/csrc/include/natten/cuda/fmha_hopper/device/fmha_sm90.hpp new file mode 100644 index 00000000..1d208f06 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_hopper/device/fmha_sm90.hpp @@ -0,0 +1,288 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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 + +#include +// common +#include "cutlass/cutlass.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 Sm90DeviceKernel { + 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("Sm90DeviceKernel::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_sm90, + 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_sm90, + 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( + "Sm90DeviceKernel::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_sm90, + 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( + "Sm90DeviceKernel()::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) { + CUTLASS_TRACE_HOST("Sm90DeviceKernel::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_sm90; + void* kernel_params[] = {¶ms}; + launch_result = ClusterLauncher::launch( + grid, cluster, block, smem_size, stream, kernel, kernel_params); + } else { + launch_result = Status::kSuccess; + cutlass::arch::synclog_setup(); + device_kernel_sm90<<>>(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/natten/csrc/include/natten/cuda/fmha_hopper/fmha_backward.cuh b/natten/csrc/include/natten/cuda/fmha_hopper/fmha_backward.cuh new file mode 100644 index 00000000..cd23dfaf --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_hopper/fmha_backward.cuh @@ -0,0 +1,224 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ + +#pragma once + +#include "cute/tensor.hpp" + +#include "cutlass/cutlass.h" +#include "cutlass/kernel_hardware_info.h" + +#include "natten/cuda/fmha_hopper/collective/fmha_fusion.hpp" +#include "natten/cuda/fmha_hopper/device/fmha_bwd_sm90.hpp" + +namespace natten { +namespace cuda { +namespace fmha_hopper { + +#if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) + +using namespace cute; +using namespace cutlass::fmha; + +template +struct KernelBackward { + using ElementAccumulator = float; + using VariableLength = cutlass::fmha::collective::VariableLength; + + // B H Q K D + using ProblemShapeRegular = cute::tuple; + using ProblemShapeVarlen = + cute::tuple; + using ProblemShapeType = + std::conditional_t; + + using Operation = cutlass::fmha::device::FmhaBwdSm90< + ProblemShapeType, + Element, + ElementAccumulator, + TileShape, + Mask>; + + using Arguments = typename Operation::Arguments; + + Operation op; + + Arguments initialize( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch, + int seqlen_Q, + int seqlen_KV, + int heads, + int dim, + float attn_scale, + // varlen parameters + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id) { + auto dim_aligned = cutlass::round_up(dim, 8); // alignment + + auto problem_shape_regular = ProblemShapeRegular{ + batch, + heads, + seqlen_Q, + seqlen_KV, + dim_aligned, + }; + + ProblemShapeType problem_shape_launch; + decltype(problem_shape_regular) problem_shape_memory; + + if constexpr (kIsVarlen) { + problem_shape_memory = problem_shape_regular; + get<0>(problem_shape_memory) = 1; + + get<0>(problem_shape_launch) = get<0>(problem_shape_regular); + get<1>(problem_shape_launch) = get<1>(problem_shape_regular); + + get<2>(problem_shape_launch) = VariableLength{ + max_seqlen_Q, + reinterpret_cast(ptr_cumulative_seqlen_Q), + seqlen_Q}; + get<3>(problem_shape_launch) = VariableLength{ + max_seqlen_KV, + reinterpret_cast(ptr_cumulative_seqlen_KV), + seqlen_KV}; + get<4>(problem_shape_launch) = get<4>(problem_shape_regular); + } else { + problem_shape_memory = problem_shape_regular; + problem_shape_launch = problem_shape_regular; + } + + int B = size<0>(problem_shape_memory); + int H = size<1>(problem_shape_memory); + int SQ = size<2>(problem_shape_memory); + int SK = size<3>(problem_shape_memory); + int D = size<4>(problem_shape_memory); + + // heads last profile, with torch's "contiguous layout" + // shape: (batch, heads, seqlen, dim) + // stride: (dim*heads*seqlen, dim*heads, dim, 1) + auto stride_Q = make_stride( + B == 1 ? 0 + : static_cast(dim_aligned * heads) * + static_cast(seqlen_Q), + dim_aligned, + dim_aligned * heads, + _1{}); + auto stride_O = stride_Q; + auto stride_K = make_stride( + B == 1 ? 0 + : static_cast(dim_aligned * heads) * + static_cast(seqlen_KV), + dim_aligned, + dim_aligned * heads, + _1{}); + auto stride_V = stride_K; + auto stride_LSE = make_stride( + B == 1 ? 0 + : static_cast(heads) * static_cast(seqlen_Q), + _1{}, + heads); + + cutlass::KernelHardwareInfo hw_info; + hw_info.device_id = device_id; + hw_info.sm_count = + cutlass::KernelHardwareInfo::query_device_multiprocessor_count( + hw_info.device_id); + + Arguments arguments{ + problem_shape_launch, + reinterpret_cast(ptr_Q), + stride_Q, + reinterpret_cast(ptr_K), + stride_K, + reinterpret_cast(ptr_V), + stride_V, + reinterpret_cast(ptr_O), + stride_O, + reinterpret_cast(ptr_LSE), + stride_LSE, + reinterpret_cast(ptr_dO), + stride_O, + reinterpret_cast(ptr_dQ), + stride_Q, + reinterpret_cast(ptr_dK), + stride_K, + reinterpret_cast(ptr_dV), + stride_V, + attn_scale, + hw_info}; + + return arguments; + } + + size_t get_workspace_size(Arguments const& arguments) { + return Operation::get_workspace_size(arguments); + } + + void run( + Arguments const& arguments, + void* workspace_ptr, + cudaStream_t stream) { + 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; + return; + } + + status = op.initialize(arguments, workspace_ptr, stream); + if (status != cutlass::Status::kSuccess) { + std::cerr + << "Failed to initialize the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(cudaGetLastError()) << std::endl; + return; + } + + // 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; + } + } +}; + +#endif // defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) + +} // namespace fmha_hopper +} // namespace cuda +} // namespace natten diff --git a/natten/csrc/include/natten/cuda/fmha_hopper/fmha_forward.cuh b/natten/csrc/include/natten/cuda/fmha_hopper/fmha_forward.cuh new file mode 100644 index 00000000..fbb0775d --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_hopper/fmha_forward.cuh @@ -0,0 +1,269 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ + +#pragma once + +#include "cute/tensor.hpp" + +#include "cutlass/cutlass.h" +#include "cutlass/kernel_hardware_info.h" + +#include "natten/cuda/fmha_hopper/collective/fmha_fusion.hpp" +#include "natten/cuda/fmha_hopper/device/fmha_sm90.hpp" +#include "natten/cuda/fmha_hopper/kernel/fmha_kernel_builder.hpp" + +#include "natten/cuda/hopper_fmha_fna.h" + +namespace natten { +namespace cuda { +namespace fmha_hopper { + +#if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) + +using namespace cute; +using namespace cutlass::fmha::kernel; +using namespace cutlass::fmha::collective; +using namespace cutlass::fmha; + +template < + typename Element, + class TileShape, + natten::cuda::hopper::HopperKernelSchedule KernelSchedule, + class Mask, + bool kIsVarlen> +struct KernelForward { + using ElementAccumulatorQK = float; + using ElementAccumulatorPV = float; + using ElementOut = Element; + using VariableLength = cutlass::fmha::collective::VariableLength; + + // B H Q K D + using ProblemShapeRegular = cute::tuple; + using ProblemShapeVarlen = + cute::tuple; + using ProblemShapeType = + std::conditional_t; + + using StrideQ = cute::tuple>; // Q D (B H) + using StrideK = cute::tuple>; // K D (B H) + + using StrideV = + StrideK; // NOTE: StrideV is different for FP8 due to transpose + using StrideO = StrideQ; + + using StrideLSE = cute::tuple>; // Q (B H) + + static_assert( + KernelSchedule == + natten::cuda::hopper::HopperKernelSchedule::NonPersistent || + KernelSchedule == + natten::cuda::hopper::HopperKernelSchedule::WSCooperative || + KernelSchedule == natten::cuda::hopper::HopperKernelSchedule::WSPingpong); + + using DispatchPolicy = std::conditional_t< + KernelSchedule == + natten::cuda::hopper::HopperKernelSchedule::NonPersistent, + cutlass::gemm::KernelTma, + std::conditional_t< + KernelSchedule == + natten::cuda::hopper::HopperKernelSchedule::WSCooperative, + cutlass::gemm::KernelTmaWarpSpecializedCooperative, + cutlass::gemm::KernelTmaWarpSpecializedPingpong>>; + + using Operation = cutlass::fmha::device::Sm90DeviceKernel< + typename cutlass::fmha::kernel::FmhaBuilder< + ProblemShapeType, + Element, + ElementAccumulatorQK, + ElementAccumulatorPV, + TileShape, + StrideQ, + StrideK, + StrideV, + Mask, + DispatchPolicy>::Kernel>; + + using Arguments = typename Operation::Arguments; + + Operation op; + + Arguments initialize( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch, + int seqlen_Q, + int seqlen_KV, + int heads, + int dim, + float attn_scale, + // varlen parameters + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id) { + auto dim_aligned = cutlass::round_up(dim, 8); // alignment + + auto problem_shape_regular = ProblemShapeRegular{ + batch, + heads, + seqlen_Q, + seqlen_KV, + dim_aligned, + }; + + ProblemShapeType problem_shape_launch; + decltype(problem_shape_regular) problem_shape_memory; + + if constexpr (kIsVarlen) { + problem_shape_memory = problem_shape_regular; + get<0>(problem_shape_memory) = 1; + + get<0>(problem_shape_launch) = get<0>(problem_shape_regular); + get<1>(problem_shape_launch) = get<1>(problem_shape_regular); + + get<2>(problem_shape_launch) = VariableLength{ + max_seqlen_Q, + reinterpret_cast(ptr_cumulative_seqlen_Q), + seqlen_Q}; + get<3>(problem_shape_launch) = VariableLength{ + max_seqlen_KV, + reinterpret_cast(ptr_cumulative_seqlen_KV), + seqlen_KV}; + get<4>(problem_shape_launch) = get<4>(problem_shape_regular); + } else { + problem_shape_memory = problem_shape_regular; + problem_shape_launch = problem_shape_regular; + } + + int B = size<0>(problem_shape_memory); + int H = size<1>(problem_shape_memory); + int SQ = size<2>(problem_shape_memory); + int SK = size<3>(problem_shape_memory); + int D = size<4>(problem_shape_memory); + + // heads last profile, with torch's "contiguous layout" + // shape: (batch, heads, seqlen, dim) + // stride: (dim*heads*seqlen, dim*heads, dim, 1) + auto stride_Q = make_stride( + heads * dim_aligned, + _1{}, + make_stride( + B == 1 ? 0 + : static_cast(heads * dim_aligned) * + static_cast(seqlen_Q), + dim_aligned)); + auto stride_O = stride_Q; + auto stride_K = make_stride( + heads * dim_aligned, + _1{}, + make_stride( + B == 1 ? 0 + : static_cast(heads * dim_aligned) * + static_cast(seqlen_KV), + dim_aligned)); + auto stride_V = stride_K; + auto stride_LSE = make_stride( + heads, + make_stride( + B == 1 + ? 0 + : static_cast(heads) * static_cast(seqlen_Q), + _1{})); + + cutlass::KernelHardwareInfo hw_info; + hw_info.device_id = device_id; + hw_info.sm_count = + cutlass::KernelHardwareInfo::query_device_multiprocessor_count( + hw_info.device_id); + + Arguments arguments{ + problem_shape_launch, + {reinterpret_cast(ptr_Q), + stride_Q, + reinterpret_cast(ptr_K), + stride_K, + reinterpret_cast(ptr_V), + stride_V, + attn_scale}, + {reinterpret_cast(ptr_O), + stride_O, + reinterpret_cast(ptr_LSE), + stride_LSE}, + hw_info}; + + return arguments; + } + + size_t get_workspace_size(Arguments const& arguments) { + return Operation::get_workspace_size(arguments); + } + + void run( + Arguments const& arguments, + void* workspace_ptr, + cudaStream_t stream) { + 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; + return; + } + + status = op.initialize(arguments, workspace_ptr, stream); + if (status != cutlass::Status::kSuccess) { + std::cerr + << "Failed to initialize the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(cudaGetLastError()) << std::endl; + return; + } + + // 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; + } + +#if 0 + cudaError_t result = cudaDeviceSynchronize(); + if (result != cudaSuccess) { + std::cerr << "Error running the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(result) << std::endl; + return; + } +#endif + } +}; + +#endif // defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) + +} // namespace fmha_hopper +} // namespace cuda +} // namespace natten diff --git a/natten/csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_builder.hpp b/natten/csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_builder.hpp new file mode 100644 index 00000000..1a4d9fd6 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_builder.hpp @@ -0,0 +1,195 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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 "natten/cuda/fmha_hopper/collective/fmha_collective_tma.hpp" +#include "natten/cuda/fmha_hopper/collective/fmha_collective_tma_warpspecialized.hpp" +#include "natten/cuda/fmha_hopper/collective/fmha_epilogue.hpp" +#include "natten/cuda/fmha_hopper/collective/fmha_varlen.hpp" +#include "natten/cuda/fmha_hopper/kernel/fmha_kernel_tma.hpp" +#include "natten/cuda/fmha_hopper/kernel/fmha_kernel_tma_warpspecialized.hpp" +#include "natten/cuda/fmha_hopper/kernel/fmha_options.hpp" + +namespace cutlass::fmha::kernel { + +template < + class ProblemShape, + class Element_, + class ElementAccumulatorQK_, + class ElementAccumulatorPV_, + class TileShape_, // BlockQO, BlockKV, BlockHead + class LayoutQ_, + class LayoutK_, + class LayoutV_, + class Fusion, + class DispatchPolicy, + class... Options> +struct FmhaBuilder; + +template < + class ProblemShape, + class Element, + class ElementAccumulator, + class TileShape, // BlockQO, BlockKV, BlockHead + class Fusion, + class... Options> +struct FmhaBuilder< + ProblemShape, + Element, + ElementAccumulator, + ElementAccumulator, + TileShape, + cute::tuple>, + cute::tuple>, + cute::tuple>, + Fusion, + cutlass::gemm::KernelTma, + Options...> { + using CollectiveMainloop = cutlass::fmha::collective::FmhaMainloopTmaSm90< + Element, + ElementAccumulator, + TileShape, + Fusion, + Options...>; + + using CollectiveEpilogue = cutlass::fmha::collective::FmhaFwdEpilogueSm90< + Element, + ElementAccumulator, + typename CollectiveMainloop::TileShapePV, + cutlass::fmha::collective::is_problem_shape_variable_length( + ProblemShape{})>; + + using Kernel = cutlass::fmha::kernel::FmhaKernelTma< + ProblemShape, + CollectiveMainloop, + CollectiveEpilogue, + Options...>; +}; + +template < + class ProblemShape, + class Element, + class ElementAccumulatorQK, + class ElementAccumulatorPV, + class TileShape, // BlockQO, BlockKV, BlockHead + class LayoutQ, + class LayoutK, + class LayoutV, + class Fusion, + class... Options> +struct FmhaBuilder< + ProblemShape, + Element, + ElementAccumulatorQK, + ElementAccumulatorPV, + TileShape, + LayoutQ, + LayoutK, + LayoutV, + Fusion, + cutlass::gemm::KernelTmaWarpSpecializedCooperative, + Options...> { + using CollectiveMainloop = + cutlass::fmha::collective::FmhaMainloopTmaWarpSpecializedSm90< + Element, + ElementAccumulatorQK, + ElementAccumulatorPV, + TileShape, + LayoutQ, + LayoutK, + LayoutV, + Fusion, + Options...>; + + using CollectiveEpilogue = cutlass::fmha::collective::FmhaFwdEpilogueSm90< + Element, + ElementAccumulatorPV, + typename CollectiveMainloop::TileShapePV, + cutlass::fmha::collective::is_problem_shape_variable_length( + ProblemShape{})>; + + static constexpr bool kIsPersistent = + find_option_t::value; + using TileScheduler = std::conditional_t< + kIsPersistent, + cutlass::fmha::kernel::PersistentTileScheduler, + cutlass::fmha::kernel::IndividualTileScheduler>; + + using Kernel = cutlass::fmha::kernel::FmhaKernelTmaWarpSpecialized< + ProblemShape, + CollectiveMainloop, + CollectiveEpilogue, + TileScheduler, + Options...>; +}; + +template < + class ProblemShape, + class Element, + class ElementAccumulatorQK, + class ElementAccumulatorPV, + class TileShape, // BlockQO, BlockKV, BlockHead + class LayoutQ, + class LayoutK, + class LayoutV, + class Fusion, + class... Options> +struct FmhaBuilder< + ProblemShape, + Element, + ElementAccumulatorQK, + ElementAccumulatorPV, + TileShape, + LayoutQ, + LayoutK, + LayoutV, + Fusion, + cutlass::gemm::KernelTmaWarpSpecializedPingpong, + Options...> { + using Kernel = typename FmhaBuilder< + ProblemShape, + Element, + ElementAccumulatorQK, + ElementAccumulatorPV, + TileShape, + LayoutQ, + LayoutK, + LayoutV, + Fusion, + cutlass::gemm::KernelTmaWarpSpecializedCooperative, + Options..., + Option, + Option>::Kernel; +}; + +} // namespace cutlass::fmha::kernel diff --git a/natten/csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_bwd_convert.hpp b/natten/csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_bwd_convert.hpp new file mode 100644 index 00000000..2325c878 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_bwd_convert.hpp @@ -0,0 +1,210 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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 "natten/cuda/fmha_hopper/collective/fmha_varlen.hpp" + +namespace cutlass::fmha::kernel { + +using namespace cute; +using namespace cutlass::fmha::collective; + +template +struct FmhaKernelBwdConvert { + struct Arguments { + ProblemShape problem_size; + + const ElementAccumulator* ptr_src_dQ; + tuple stride_src_dQ; + const ElementAccumulator* ptr_src_dK; + tuple stride_src_dK; + const ElementAccumulator* ptr_src_dV; + tuple stride_src_dV; + + Element* ptr_dest_dQ; + tuple stride_dest_dQ; + Element* ptr_dest_dK; + tuple stride_dest_dK; + Element* ptr_dest_dV; + tuple stride_dest_dV; + }; + + using Params = Arguments; + + using ClusterShape = Shape<_1, _1, _1>; + static constexpr int SharedStorageSize = 0; + + static const int MinBlocksPerMultiprocessor = 1; + static const int MaxThreadsPerBlock = 128; + using ArchTag = cutlass::arch::Sm90; + + static const int kBlockSeq = 8; + + 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 const int kNumThreadsD = 16; + static const int kNumThreadsSeq = MaxThreadsPerBlock / kNumThreadsD; + static const int kElementsPerLoad = 4; + + static const int kIterationsSeq = kBlockSeq / kNumThreadsSeq; + + static bool can_implement(Arguments const& args) { + return get<4>(args.problem_size) % kElementsPerLoad == 0; + } + + static dim3 get_grid_shape(Params const& params) { + dim3 grid( + // never put seq in z, long seqs can easily exceed the 64K limit + ceil_div( + std::max( + size<2>(params.problem_size), size<3>(params.problem_size)), + kBlockSeq), + size<1>(params.problem_size), + size<0>(params.problem_size)); + return grid; + } + + static dim3 get_block_shape() { + dim3 block(kNumThreadsD, kNumThreadsSeq, 1); + return block; + } + + static Params to_underlying_arguments( + Arguments const& args, + void* workspace) { + return args; + } + + template + CUTLASS_DEVICE void copy( + Params const& params, + const ElementAccumulator* ptr_src, + StrideSrc const& stride_src, + Element* ptr_dest, + StrideDest const& stride_dest, + Count const& count) { + auto ptr_src_bh = ptr_src + + static_cast(get<0>(stride_src)) * + static_cast(blockIdx.z) + + static_cast(get<1>(stride_src)) * + static_cast(blockIdx.y); + auto ptr_dest_bh = ptr_dest + + static_cast(get<0>(stride_dest)) * + static_cast(blockIdx.z) + + static_cast(get<1>(stride_dest)) * + static_cast(blockIdx.y); + int seqlen = count; + if constexpr (is_variable_length_v) { + int offset = count.cumulative_length[blockIdx.z]; + ptr_dest_bh += static_cast(offset) * + static_cast(get<2>(stride_dest)); + seqlen = count.cumulative_length[blockIdx.z + 1] - offset; + } + + for (int idx_s_t = threadIdx.y; idx_s_t < kBlockSeq; + idx_s_t += kNumThreadsSeq) { + int idx_s = idx_s_t + kBlockSeq * blockIdx.x; + if (idx_s >= seqlen) + continue; + auto ptr_src_bhs = ptr_src_bh + + static_cast(idx_s) * + static_cast(get<2>(stride_src)); + auto ptr_dest_bhs = ptr_dest_bh + + static_cast(idx_s) * + static_cast(get<2>(stride_dest)); + + for (int idx_d = threadIdx.x * kElementsPerLoad; + idx_d < get<4>(params.problem_size); + idx_d += kElementsPerLoad * kNumThreadsD) { + ElementAccumulator value_src[kElementsPerLoad]; + Element value_dest[kElementsPerLoad]; + + using VecSrc = + uint_bit_t * kElementsPerLoad>; + using VecDest = uint_bit_t * kElementsPerLoad>; + *reinterpret_cast(value_src) = + *reinterpret_cast(&ptr_src_bhs[idx_d]); + + for (int v = 0; v < kElementsPerLoad; v++) { + value_dest[v] = static_cast(value_src[v]); + } + + *reinterpret_cast(&ptr_dest_bhs[idx_d]) = + *reinterpret_cast(value_dest); + } + } + } + + CUTLASS_DEVICE void operator()(const Params& params, char* smem) { + if (params.ptr_src_dQ != nullptr) { + copy( + params, + params.ptr_src_dQ, + params.stride_src_dQ, + params.ptr_dest_dQ, + params.stride_dest_dQ, + get<2>(params.problem_size)); + } + if (params.ptr_src_dK != nullptr) { + copy( + params, + params.ptr_src_dK, + params.stride_src_dK, + params.ptr_dest_dK, + params.stride_dest_dK, + get<3>(params.problem_size)); + } + if (params.ptr_src_dV != nullptr) { + copy( + params, + params.ptr_src_dV, + params.stride_src_dV, + params.ptr_dest_dV, + params.stride_dest_dV, + get<3>(params.problem_size)); + } + } +}; + +} // namespace cutlass::fmha::kernel diff --git a/natten/csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_bwd_sum_OdO.hpp b/natten/csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_bwd_sum_OdO.hpp new file mode 100644 index 00000000..2edebdb6 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_bwd_sum_OdO.hpp @@ -0,0 +1,192 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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 "natten/cuda/fmha_hopper/collective/fmha_varlen.hpp" + +namespace cutlass::fmha::kernel { + +using namespace cute; +using namespace cutlass::fmha::collective; + +template +struct FmhaKernelBwdSumOdO { + struct Arguments { + ProblemShape problem_shape; + + const Element* ptr_O; + cute::tuple stride_O; + const Element* ptr_dO; + cute::tuple stride_dO; + + ElementAcc* ptr_sum_OdO; + cute::tuple stride_sum_OdO; + + const ElementAcc* ptr_lse; + // NATTEN has a different LSE layout + cute::tuple stride_lse; + + ElementAcc* ptr_scaled_lse; + cute::tuple stride_scaled_lse; + + // NOTE (ahassani): I tried doing the lse scale here the same way it's done + // in blackwell fmha, but for some reason that I don't have time to + // investigate, this causes some minor instability on dQ allcloses.... I + // don't want to assume it's nothing to worry about and sweep it under the + // rug, so rolling it back. + // ElementAcc sum_odo_scale = 1.0; + // ElementAcc lse_scale = 1.0; + }; + + using Params = Arguments; + + using ClusterShape = Shape<_1, _1, _1>; + static constexpr int SharedStorageSize = 0; + + static const int MinBlocksPerMultiprocessor = 1; + static const int MaxThreadsPerBlock = 128; + using ArchTag = cutlass::arch::Sm90; + + 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 const int kBlockQ = 16; + + static const int kNumThreadsD = 8; + static const int kNumThreadsQ = MaxThreadsPerBlock / kNumThreadsD; + static const int kElementsPerLoad = 2; + + static const int kIterationsQ = kBlockQ / kNumThreadsQ; + + static bool can_implement(Arguments const& args) { + return get<4>(args.problem_shape) % kElementsPerLoad == 0; + } + + static dim3 get_grid_shape(Params const& params) { + dim3 grid( + ceil_div(size<2>(params.problem_shape), kBlockQ), + size<1>(params.problem_shape), + size<0>(params.problem_shape)); + return grid; + } + + static dim3 get_block_shape() { + dim3 block(kNumThreadsD, kNumThreadsQ, 1); + return block; + } + + static Params to_underlying_arguments( + Arguments const& args, + void* workspace) { + return args; + } + + CUTLASS_DEVICE void operator()(const Params& params, char* smem) { + auto ptr_O_bh = params.ptr_O + blockIdx.y * get<1>(params.stride_O) + + blockIdx.z * get<0>(params.stride_O); + auto ptr_dO_bh = params.ptr_dO + blockIdx.y * get<1>(params.stride_dO) + + blockIdx.z * get<0>(params.stride_dO); + auto ptr_sum_OdO_bh = params.ptr_sum_OdO + + blockIdx.y * get<1>(params.stride_sum_OdO) + + blockIdx.z * get<0>(params.stride_sum_OdO); + auto ptr_lse_bh = params.ptr_lse + blockIdx.y * get<1>(params.stride_lse) + + blockIdx.z * get<0>(params.stride_lse); + auto ptr_scaled_lse_bh = params.ptr_scaled_lse + + blockIdx.y * get<1>(params.stride_scaled_lse) + + blockIdx.z * get<0>(params.stride_scaled_lse); + + auto problem_q = get<2>(params.problem_shape); + int seqlen_q = problem_q; + if constexpr (is_problem_shape_variable_length(ProblemShape{})) { + int offset = problem_q.cumulative_length[blockIdx.z]; + ptr_O_bh += offset * get<2>(params.stride_O); + ptr_dO_bh += offset * get<2>(params.stride_dO); + ptr_lse_bh += offset * get<2>(params.stride_lse); + seqlen_q = problem_q.cumulative_length[blockIdx.z + 1] - offset; + } + + CUTLASS_PRAGMA_UNROLL + for (int idx_q_t = threadIdx.y; idx_q_t < kBlockQ; + idx_q_t += kNumThreadsQ) { + int idx_q = idx_q_t + kBlockQ * blockIdx.x; + if (idx_q >= seqlen_q) + continue; + ElementAcc acc = 0; + auto ptr_O_bhq = ptr_O_bh + idx_q * get<2>(params.stride_O); + auto ptr_dO_bhq = ptr_dO_bh + idx_q * get<2>(params.stride_dO); + auto ptr_sum_OdO_bhq = + ptr_sum_OdO_bh + idx_q * get<2>(params.stride_sum_OdO); + auto ptr_lse_bhq = ptr_lse_bh + idx_q * get<2>(params.stride_lse); + auto ptr_scaled_lse_bhq = + ptr_scaled_lse_bh + idx_q * get<2>(params.stride_scaled_lse); + + for (int idx_d = threadIdx.x * kElementsPerLoad; + idx_d < get<4>(params.problem_shape); + idx_d += kElementsPerLoad * kNumThreadsD) { + Element value_O[kElementsPerLoad]; + Element value_dO[kElementsPerLoad]; + + using Vec = uint_bit_t * kElementsPerLoad>; + *reinterpret_cast(value_O) = + *reinterpret_cast(&ptr_O_bhq[idx_d]); + *reinterpret_cast(value_dO) = + *reinterpret_cast(&ptr_dO_bhq[idx_d]); + + for (int v = 0; v < kElementsPerLoad; v++) { + acc += value_O[v] * value_dO[v]; + } + } + + for (int i = 1; i < kNumThreadsD; i *= 2) { + acc += __shfl_xor_sync((uint32_t)-1, acc, i, kNumThreadsD); + } + + if (threadIdx.x == 0) { + *ptr_sum_OdO_bhq = /*params.sum_odo_scale * */ acc; + *ptr_scaled_lse_bhq = /*params.lse_scale * */ *ptr_lse_bhq; + } + } + } +}; + +} // namespace cutlass::fmha::kernel diff --git a/natten/csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_tma.hpp b/natten/csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_tma.hpp new file mode 100644 index 00000000..1250bef1 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_tma.hpp @@ -0,0 +1,255 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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/arch/arch.h" +#include "cutlass/cutlass.h" +#include "cutlass/pipeline/pipeline.hpp" + +#include "natten/cuda/fmha_hopper/kernel/fmha_options.hpp" +#include "natten/cuda/fmha_hopper/kernel/fmha_tile_scheduler.hpp" + +namespace cutlass::fmha::kernel { + +template < + class ProblemShape, + class CollectiveMainloop, + class CollectiveEpilogue, + class... Options> +struct FmhaKernelTma { + // Options + static constexpr int kBlocksPerSM = + find_option_t, Options...>::value; + + using Element = typename CollectiveMainloop::Element; + using ElementAccumulator = typename CollectiveMainloop::ElementAccumulator; + + using TileScheduler = IndividualTileScheduler; + + using StagesQ = typename CollectiveMainloop::StagesQ; + using Stages = typename CollectiveMainloop::Stages; + + using TileShape = typename CollectiveMainloop::TileShape; + using ClusterShape = typename CollectiveMainloop::ClusterShape; + + using MainloopPipeline = typename CollectiveMainloop::MainloopPipeline; + using MainloopPipelineQ = typename CollectiveMainloop::MainloopPipelineQ; + + using SmemLayoutQ = typename CollectiveMainloop::SmemLayoutQ; + using SmemLayoutK = typename CollectiveMainloop::SmemLayoutK; + + struct SharedStorage { + union { + typename CollectiveMainloop::SharedStorage mainloop; + typename CollectiveEpilogue::TensorStorage epilogue; + }; + + using PipelineStorage = typename MainloopPipeline::SharedStorage; + using PipelineStorageQ = typename MainloopPipelineQ::SharedStorage; + alignas(16) PipelineStorage pipeline_storage; + alignas(16) PipelineStorageQ pipeline_storage_q; + + using EpiLoadPipelineStorage = typename CollectiveEpilogue::PipelineStorage; + alignas(16) EpiLoadPipelineStorage epi_load; + }; + + static constexpr int SharedStorageSize = sizeof(SharedStorage); + + struct Arguments { + ProblemShape problem_size; + typename CollectiveMainloop::Arguments mainloop; + typename CollectiveEpilogue::Arguments epilogue; + KernelHardwareInfo hw_info; + }; + + struct Params { + ProblemShape problem_size; + typename CollectiveMainloop::Params mainloop; + typename CollectiveEpilogue::Params epilogue; + typename TileScheduler::Params tile_scheduler; + }; + + using PipelineParams = typename MainloopPipeline::Params; + using PipelineState = + typename cutlass::PipelineState; + using PipelineParamsQ = typename MainloopPipelineQ::Params; + using PipelineStateQ = + typename cutlass::PipelineState; + + static const int MinBlocksPerMultiprocessor = kBlocksPerSM; + static const int MaxThreadsPerBlock = CollectiveMainloop::MaxThreadsPerBlock; + using ArchTag = cutlass::arch::Sm90; + + 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_size, 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_size, + CollectiveMainloop::to_underlying_arguments( + args.problem_size, args.mainloop, workspace), + CollectiveEpilogue::to_underlying_arguments( + args.problem_size, args.epilogue, workspace), + TileScheduler::to_underlying_arguments( + args.problem_size, args.hw_info, ClusterShape{}, TileShape{})}; + } + + CUTLASS_DEVICE void operator()(const Params& params, char* smem) { +#if !defined(CUTLASS_ARCH_MMA_SM90A_ENABLED) + CUTE_INVALID_CONTROL_PATH( + "ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); +#else + TileScheduler tile_scheduler{params.tile_scheduler}; + + // Shared memory. + auto& storage = *reinterpret_cast(smem); + + int thread_idx = int(threadIdx.x); + + uint32_t block_rank_in_cluster = cute::block_rank_in_cluster(); + + int warp_idx = cutlass::canonical_warp_idx_sync(); + int warp_group_thread_idx = thread_idx % cutlass::NumThreadsPerWarpGroup; + int lane_predicate = cute::elect_one_sync(); + + // Issue Tma Descriptor Prefetch from a single thread + if ((warp_idx == 0) && lane_predicate) { + CollectiveMainloop::prefetch_tma_descriptors(params.mainloop); + } + + PipelineParamsQ pipeline_params_q; + pipeline_params_q.transaction_bytes = + size(SmemLayoutQ{}(_, _, _0{})) * sizeof(Element); // Q + pipeline_params_q.role = + MainloopPipelineQ::ThreadCategory::ProducerConsumer; + pipeline_params_q.is_leader = warp_group_thread_idx == 0; + pipeline_params_q.num_consumers = cutlass::NumThreadsPerWarpGroup; + + PipelineParams pipeline_params; + pipeline_params.transaction_bytes = + size(SmemLayoutK{}(_, _, _0{})) * sizeof(Element); // KV + pipeline_params.role = MainloopPipeline::ThreadCategory::ProducerConsumer; + pipeline_params.is_leader = warp_group_thread_idx == 0; + pipeline_params.num_consumers = cutlass::NumThreadsPerWarpGroup; + + MainloopPipelineQ pipeline_q( + storage.pipeline_storage_q, pipeline_params_q, Shape<_1, _1, _1>{}); + MainloopPipeline pipeline( + storage.pipeline_storage, pipeline_params, ClusterShape{}); + + using EpiLoadPipeline = typename CollectiveEpilogue::LoadPipeline; + typename EpiLoadPipeline::Params epi_load_pipeline_params; + epi_load_pipeline_params.role = + EpiLoadPipeline::ThreadCategory::ProducerConsumer; + epi_load_pipeline_params.dst_blockid = cute::block_rank_in_cluster(); + epi_load_pipeline_params.producer_arv_count = NumThreadsPerWarp; + epi_load_pipeline_params.consumer_arv_count = NumThreadsPerWarpGroup; + epi_load_pipeline_params.transaction_bytes = + CollectiveEpilogue::TmaTransactionBytes; + EpiLoadPipeline epi_load_pipeline( + storage.epi_load, epi_load_pipeline_params); + + // State variables used for iterating the circular buffer + // smem_pipe_read / release is used by the consumer of SMEM data - i.e MMA + // smem_pipe_write is used by the producer of SMEM data - i.e TMA + PipelineState smem_pipe_read; + PipelineState smem_pipe_write = + cutlass::make_producer_start_state(); + + PipelineStateQ smem_pipe_read_q; + PipelineStateQ smem_pipe_write_q = + cutlass::make_producer_start_state(); + + // We need this to guarantee that the Pipeline init is visible + // To all producers and consumer blocks in the Cluster + // and to finish smem init + if constexpr (size(ClusterShape{}) > 1) { + cute::cluster_arrive_relaxed(); + cute::cluster_wait(); + } else { + __syncthreads(); + } + + auto blk_coord = tile_scheduler.get_block_coord(); + + CollectiveMainloop collective_mainloop; + auto result = collective_mainloop.compute( + block_rank_in_cluster, + blk_coord, + params.mainloop, + params.problem_size, + pipeline, + smem_pipe_read, + smem_pipe_write, + pipeline_q, + smem_pipe_read_q, + smem_pipe_write_q, + storage.mainloop); + + CollectiveEpilogue epilogue; + epilogue( + typename CollectiveMainloop::TileShapePV{}, + blk_coord, + result, + typename CollectiveMainloop::TiledMmaPV{}, + params.problem_size, + params.epilogue, + epi_load_pipeline, + storage.epilogue); +#endif + } +}; + +} // namespace cutlass::fmha::kernel diff --git a/natten/csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_tma_warpspecialized.hpp b/natten/csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_tma_warpspecialized.hpp new file mode 100644 index 00000000..2466d046 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_tma_warpspecialized.hpp @@ -0,0 +1,491 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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/arch/arch.h" +#include "cutlass/arch/reg_reconfig.h" +#include "cutlass/cutlass.h" +#include "cutlass/pipeline/pipeline.hpp" + +#include "natten/cuda/fmha_hopper/kernel/fmha_options.hpp" + +namespace cutlass::fmha::kernel { + +using namespace cute; + +template < + class ProblemShape, + class CollectiveMainloop, + class CollectiveEpilogue, + class TileScheduler, + class... Options> +struct FmhaKernelTmaWarpSpecialized { + // Options + static constexpr bool kIsEpilogueLocked = + find_option_t::value; + static constexpr bool kLoadsQSeparately = + find_option_t::value; + + static const int NumLoadWarpGroups = 1; + static constexpr int NumMmaWarpGroups = CollectiveMainloop::NumMmaWarpGroups; + + using TileShape = typename CollectiveMainloop::TileShape; + using ClusterShape = typename CollectiveMainloop::ClusterShape; + + using MainloopPipelineOuter = typename CollectiveMainloop::MainloopPipelineQ; + using MainloopPipelineInner = typename CollectiveMainloop::MainloopPipeline; + using MainloopPipelineReducer = cutlass::PipelineAsync<2>; + + static constexpr uint32_t StagesPerMathWarpGroup = 2; + using MathWarpGroupOrderBarrier = + cutlass::OrderedSequenceBarrier; + + struct TensorStorageStruct { + typename CollectiveMainloop::SharedStorage mainloop; + typename CollectiveEpilogue::TensorStorage epilogue[NumMmaWarpGroups]; + }; + union TensorStorageUnion { + typename CollectiveMainloop::SharedStorage mainloop; + typename CollectiveEpilogue::TensorStorage epilogue[NumMmaWarpGroups]; + }; + using TensorStorage = std::conditional_t< + CollectiveMainloop::kIsPersistent, + TensorStorageStruct, + TensorStorageUnion>; + + struct SharedStorage { + TensorStorage tensors; + + using PipelineStorageInner = typename MainloopPipelineInner::SharedStorage; + using PipelineStorageOuter = typename MainloopPipelineOuter::SharedStorage; + using PipelineStorageReducer = + typename MainloopPipelineReducer::SharedStorage; + + alignas(16) PipelineStorageInner pipeline_storage_inner; + alignas(16) PipelineStorageOuter pipeline_storage_outer; + alignas(16) PipelineStorageReducer pipeline_storage_reducer; + + using MathWarpGroupOrderBarrierStorage = + typename MathWarpGroupOrderBarrier::SharedStorage; + alignas(16) MathWarpGroupOrderBarrierStorage math_wg_order; + + alignas(16) cutlass::arch::ClusterBarrier load_warp_barrier; + + using EpiLoadPipelineStorage = typename CollectiveEpilogue::PipelineStorage; + alignas(16) EpiLoadPipelineStorage epi_load; + }; + + static constexpr int SharedStorageSize = sizeof(SharedStorage); + + struct Arguments { + ProblemShape problem_size; + typename CollectiveMainloop::Arguments mainloop; + typename CollectiveEpilogue::Arguments epilogue; + KernelHardwareInfo hw_info; + }; + + struct Params { + ProblemShape problem_size; + typename CollectiveMainloop::Params mainloop; + typename CollectiveEpilogue::Params epilogue; + typename TileScheduler::Params tile_scheduler; + }; + + using PipelineParamsInner = typename MainloopPipelineInner::Params; + using PipelineStateInner = + typename cutlass::PipelineState; + using PipelineParamsOuter = typename MainloopPipelineOuter::Params; + using PipelineStateOuter = + typename cutlass::PipelineState; + using PipelineParamsReducer = typename MainloopPipelineReducer::Params; + using PipelineStateReducer = + typename cutlass::PipelineState; + + static const int MinBlocksPerMultiprocessor = 1; + static const int MaxThreadsPerBlock = + (NumMmaWarpGroups + NumLoadWarpGroups) * cutlass::NumThreadsPerWarpGroup; + using ArchTag = cutlass::arch::Sm90; + + static constexpr uint32_t LoadRegisterRequirement = 40 - 2 * 8; + static constexpr uint32_t TotalRegisterSupply = + (64 * 1024 / MaxThreadsPerBlock / MinBlocksPerMultiprocessor / 8) * 8 * + MaxThreadsPerBlock / cutlass::NumThreadsPerWarpGroup; + static constexpr uint32_t MmaRegisterRequirement = + ((TotalRegisterSupply - LoadRegisterRequirement) / NumMmaWarpGroups / 8) * + 8; + + 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_size, 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_size, + CollectiveMainloop::to_underlying_arguments( + args.problem_size, args.mainloop, workspace), + CollectiveEpilogue::to_underlying_arguments( + args.problem_size, args.epilogue, workspace), + TileScheduler::to_underlying_arguments( + args.problem_size, args.hw_info, ClusterShape{}, TileShape{})}; + } + + CUTLASS_DEVICE void operator()(const Params& params, char* smem) { +#if !defined(CUTLASS_ARCH_MMA_SM90A_ENABLED) + CUTE_INVALID_CONTROL_PATH( + "ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); +#else + + enum class WarpGroupRole { + Producer = 0, + Consumer0 = 1, + Consumer1 = 2, + Consumer2 = 3, + Consumer3 = 4, + }; + enum class ProducerWarpRole { + LoadKV = 1, + Reducer = 0, + MaybeLoadQ = 2, // is kLoadsQSeparately is true, this warp loads Q + // (otherwise warp 0 does it) + MainloopEpilogue = 3, + }; + + static constexpr ProducerWarpRole WarpRoleLoadQ = kLoadsQSeparately + ? ProducerWarpRole::MaybeLoadQ + : ProducerWarpRole::LoadKV; + + TileScheduler tile_scheduler{params.tile_scheduler}; + + // Shared memory. + auto& storage = *reinterpret_cast(smem); + + int lane_idx = cutlass::canonical_lane_idx(); + int warp_idx = cutlass::canonical_warp_idx_sync(); + int warp_idx_in_warp_group = warp_idx % cutlass::NumWarpsPerWarpGroup; + int warp_group_idx = cutlass::canonical_warp_group_idx(); + auto warp_group_role = WarpGroupRole(warp_group_idx); + auto producer_warp_role = ProducerWarpRole(warp_idx_in_warp_group); + int consumer_warp_group_idx = + warp_group_idx - (int)WarpGroupRole::Consumer0; + int lane_predicate = cute::elect_one_sync(); + uint32_t block_rank_in_cluster = cute::block_rank_in_cluster(); + + // Issue Tma Descriptor Prefetch from a single thread + if ((warp_idx == 0) && lane_predicate) { + CollectiveMainloop::prefetch_tma_descriptors(params.mainloop); + } + + PipelineParamsOuter pipeline_params_outer; + pipeline_params_outer.transaction_bytes = + CollectiveMainloop::kOuterLoadBytes; + pipeline_params_outer.is_leader = + lane_predicate && (producer_warp_role == WarpRoleLoadQ); + pipeline_params_outer.num_consumers = cutlass::NumThreadsPerWarpGroup; + + PipelineParamsInner pipeline_params_inner; + pipeline_params_inner.transaction_bytes = + CollectiveMainloop::kInnerLoadBytes; + pipeline_params_inner.is_leader = + lane_predicate && (producer_warp_role == ProducerWarpRole::LoadKV); + pipeline_params_inner.num_consumers = + NumMmaWarpGroups * cutlass::NumThreadsPerWarpGroup; + + PipelineParamsReducer pipeline_params_reducer; + pipeline_params_reducer.producer_arv_count = + NumMmaWarpGroups * cutlass::NumThreadsPerWarpGroup; + pipeline_params_reducer.consumer_arv_count = cutlass::NumThreadsPerWarp; + + using EpiLoadPipeline = typename CollectiveEpilogue::LoadPipeline; + typename EpiLoadPipeline::Params epi_load_pipeline_params; + + if (warp_group_role == WarpGroupRole::Producer && + producer_warp_role == ProducerWarpRole::MainloopEpilogue) { + epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Producer; + } + if (warp_group_role == WarpGroupRole::Producer && + producer_warp_role == ProducerWarpRole::LoadKV) { + pipeline_params_inner.role = + MainloopPipelineInner::ThreadCategory::Producer; + } + if (warp_group_role == WarpGroupRole::Producer && + producer_warp_role == WarpRoleLoadQ) { + pipeline_params_outer.role = + MainloopPipelineOuter::ThreadCategory::Producer; + } + if (warp_group_role == WarpGroupRole::Producer && + producer_warp_role == ProducerWarpRole::Reducer) { + pipeline_params_reducer.role = + MainloopPipelineReducer::ThreadCategory::Consumer; + } + if (warp_group_role == WarpGroupRole::Consumer0 || + warp_group_role == WarpGroupRole::Consumer1 || + warp_group_role == WarpGroupRole::Consumer2 || + warp_group_role == WarpGroupRole::Consumer3) { + pipeline_params_inner.role = + MainloopPipelineInner::ThreadCategory::Consumer; + pipeline_params_outer.role = + MainloopPipelineOuter::ThreadCategory::Consumer; + pipeline_params_reducer.role = + MainloopPipelineReducer::ThreadCategory::Producer; + epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Consumer; + } + + MainloopPipelineOuter pipeline_outer( + storage.pipeline_storage_outer, + pipeline_params_outer, + Shape<_1, _1, _1>{}); + MainloopPipelineInner pipeline_inner( + storage.pipeline_storage_inner, pipeline_params_inner, ClusterShape{}); + MainloopPipelineReducer pipeline_reducer( + storage.pipeline_storage_reducer, pipeline_params_reducer); + + // State variables used for iterating the circular buffer + // smem_pipe_read / release is used by the consumer of SMEM data - i.e MMA + // smem_pipe_write is used by the producer of SMEM data - i.e TMA + PipelineStateInner smem_pipe_read_inner; + PipelineStateInner smem_pipe_write_inner = + cutlass::make_producer_start_state(); + + PipelineStateOuter smem_pipe_read_outer; + PipelineStateOuter smem_pipe_write_outer = + cutlass::make_producer_start_state(); + + PipelineStateReducer smem_pipe_read_reducer; + PipelineStateReducer smem_pipe_write_reducer = + cutlass::make_producer_start_state(); + + typename MathWarpGroupOrderBarrier::Params params_math_wg_order_barrier; + // DMA Load WG will not participate in these Ordered Barrier syncs + params_math_wg_order_barrier.group_id = consumer_warp_group_idx; + params_math_wg_order_barrier.group_size = + cutlass::NumThreadsPerWarpGroup; // Number of threads / participants in + // a group + MathWarpGroupOrderBarrier math_wg_order_barrier( + storage.math_wg_order, params_math_wg_order_barrier); + + // Epilogue Load pipeline + epi_load_pipeline_params.dst_blockid = cute::block_rank_in_cluster(); + epi_load_pipeline_params.producer_arv_count = NumThreadsPerWarp; + epi_load_pipeline_params.consumer_arv_count = NumThreadsPerWarpGroup; + epi_load_pipeline_params.transaction_bytes = + CollectiveEpilogue::TmaTransactionBytes; + EpiLoadPipeline epi_load_pipeline( + storage.epi_load, epi_load_pipeline_params); + + if constexpr (kLoadsQSeparately) { + if ((warp_idx == 0) && lane_predicate) { + storage.load_warp_barrier.init(2 * cutlass::NumThreadsPerWarp); + } + cutlass::arch::fence_barrier_init(); + } + + // We need this to guarantee that the Pipeline init is visible + // To all producers and consumer blocks in the Cluster + // and to finish smem init + if constexpr (size(ClusterShape{}) > 1) { + cute::cluster_arrive_relaxed(); + cute::cluster_wait(); + } else { + __syncthreads(); + } + + CollectiveMainloop collective_mainloop; + + if (warp_group_role == WarpGroupRole::Producer) { + cutlass::arch::warpgroup_reg_dealloc(); + if (producer_warp_role == ProducerWarpRole::LoadKV) { + bool do_barrier = kLoadsQSeparately; + + CUTLASS_PRAGMA_NO_UNROLL + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + auto blk_coord = tile_scheduler.get_block_coord(); + collective_mainloop.template load_kv_maybe_q( + block_rank_in_cluster, + blk_coord, + params.mainloop, + params.problem_size, + pipeline_inner, + smem_pipe_write_inner, + pipeline_outer, + smem_pipe_write_outer, + storage.tensors.mainloop, + storage.load_warp_barrier, + do_barrier); + do_barrier = false; + } + } else if ( + kLoadsQSeparately && + (producer_warp_role == ProducerWarpRole::MaybeLoadQ)) { + bool do_barrier = true; + + CUTLASS_PRAGMA_NO_UNROLL + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + auto blk_coord = tile_scheduler.get_block_coord(); + collective_mainloop.load_maybe_q( + blk_coord, + params.mainloop, + params.problem_size, + pipeline_outer, + smem_pipe_write_outer, + storage.tensors.mainloop, + storage.load_warp_barrier, + do_barrier); + do_barrier = false; + } + } else if (producer_warp_role == ProducerWarpRole::Reducer) { + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + auto blk_coord = tile_scheduler.get_block_coord(); + collective_mainloop.reduce( + blk_coord, + params.mainloop, + params.problem_size, + pipeline_reducer, + smem_pipe_read_reducer, + storage.tensors.mainloop); + } + } + } else if ( + warp_group_role == WarpGroupRole::Consumer0 || + warp_group_role == WarpGroupRole::Consumer1 || + warp_group_role == WarpGroupRole::Consumer2 || + warp_group_role == WarpGroupRole::Consumer3) { + cutlass::arch::warpgroup_reg_alloc(); + CUTLASS_PRAGMA_NO_UNROLL + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + auto blk_coord = tile_scheduler.get_block_coord(); + auto wg_coord = blk_coord; + + constexpr int kOuterLoads = CollectiveMainloop::kOuterLoads; + + if (warp_group_role == WarpGroupRole::Consumer0) { + smem_pipe_read_outer.advance(0 * kOuterLoads); + } else if (warp_group_role == WarpGroupRole::Consumer1) { + smem_pipe_read_outer.advance(1 * kOuterLoads); + } else if (warp_group_role == WarpGroupRole::Consumer2) { + smem_pipe_read_outer.advance(2 * kOuterLoads); + } else if (warp_group_role == WarpGroupRole::Consumer3) { + smem_pipe_read_outer.advance(3 * kOuterLoads); + } + + constexpr int wg_dim = + is_constant<0, decltype(get<1>(wg_coord))>::value ? 0 : 1; + auto& wg_block = get(wg_coord); + if (warp_group_role == WarpGroupRole::Consumer0) { + wg_block = NumMmaWarpGroups * wg_block + 0; + } else if (warp_group_role == WarpGroupRole::Consumer1) { + wg_block = NumMmaWarpGroups * wg_block + 1; + } else if (warp_group_role == WarpGroupRole::Consumer2) { + wg_block = NumMmaWarpGroups * wg_block + 2; + } else if (warp_group_role == WarpGroupRole::Consumer3) { + wg_block = NumMmaWarpGroups * wg_block + 3; + } + + auto result = collective_mainloop.compute( + blk_coord, + wg_coord, + params.mainloop, + params.problem_size, + pipeline_inner, + smem_pipe_read_inner, + pipeline_outer, + smem_pipe_read_outer, + pipeline_reducer, + smem_pipe_write_reducer, + storage.tensors.mainloop, + math_wg_order_barrier); + + if (warp_group_role == WarpGroupRole::Consumer0) { + smem_pipe_read_outer.advance(kOuterLoads * (NumMmaWarpGroups - 0)); + } + if constexpr (NumMmaWarpGroups >= 2) { + if (warp_group_role == WarpGroupRole::Consumer1) { + smem_pipe_read_outer.advance(kOuterLoads * (NumMmaWarpGroups - 1)); + } + } + if constexpr (NumMmaWarpGroups >= 3) { + if (warp_group_role == WarpGroupRole::Consumer2) { + smem_pipe_read_outer.advance(kOuterLoads * (NumMmaWarpGroups - 2)); + } + } + if constexpr (NumMmaWarpGroups >= 4) { + if (warp_group_role == WarpGroupRole::Consumer3) { + smem_pipe_read_outer.advance(kOuterLoads * (NumMmaWarpGroups - 3)); + } + } + + if constexpr (kIsEpilogueLocked) + ; + math_wg_order_barrier.wait(); + + CollectiveEpilogue epilogue; + epilogue( + typename CollectiveMainloop::TileShapePV{}, + wg_coord, + result, + typename CollectiveMainloop::TiledMmaPV{}, + params.problem_size, + params.epilogue, + epi_load_pipeline, + storage.tensors.epilogue[consumer_warp_group_idx]); + + if constexpr (kIsEpilogueLocked) + ; + math_wg_order_barrier.arrive(); + } + } +#endif + } +}; + +} // namespace cutlass::fmha::kernel diff --git a/natten/csrc/include/natten/cuda/fmha_hopper/kernel/fmha_options.hpp b/natten/csrc/include/natten/cuda/fmha_hopper/kernel/fmha_options.hpp new file mode 100644 index 00000000..4083a0c2 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_hopper/kernel/fmha_options.hpp @@ -0,0 +1,83 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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< + Option::tag == kTag, + Option, + find_option> {}; + +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/natten/csrc/include/natten/cuda/fmha_hopper/kernel/fmha_tile_scheduler.hpp b/natten/csrc/include/natten/cuda/fmha_hopper/kernel/fmha_tile_scheduler.hpp new file mode 100644 index 00000000..b127f282 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fmha_hopper/kernel/fmha_tile_scheduler.hpp @@ -0,0 +1,227 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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 IndividualTileScheduler { + struct Params { + dim3 grid; + }; + + bool valid_ = true; + + CUTLASS_DEVICE + IndividualTileScheduler(Params const&) {} + + template + static Params to_underlying_arguments( + ProblemSize const& problem_size, + KernelHardwareInfo hw_info, + ClusterShape const& cluster_shape, + TileShape const& tile_shape) { + using namespace cute; + dim3 grid( + round_up( + ceil_div(size<2>(problem_size), size<0>(tile_shape)), + size<0>(cluster_shape)), + size<0>(problem_size), + size<1>(problem_size)); + 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; + // NOTE (ahassani): we need to typecast block idx to a signed integer here, + // otherwise the non-TMA epilogue predication can fail (negatives will be + // interpreted as 0s by the compiler) + return make_coord( + static_cast(blockIdx.x), _0{}, make_coord(blockIdx.y, blockIdx.z)); + } + + CUTLASS_DEVICE + IndividualTileScheduler& operator++() { + valid_ = false; + return *this; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +struct PersistentTileScheduler { + struct Params { + int num_blocks; + FastDivmod divmod_m_block; + FastDivmod divmod_b; + FastDivmod divmod_h; + + KernelHardwareInfo hw_info; + }; + + int block_idx = 0; + Params params; + + CUTLASS_DEVICE + PersistentTileScheduler(Params const& params) + : block_idx(blockIdx.x), params(params) {} + + template + static Params to_underlying_arguments( + ProblemSize const& problem_size, + KernelHardwareInfo hw_info, + ClusterShape const& cluster_shape, + TileShape const& tile_shape) { + using namespace cute; + // Get SM count if needed, otherwise use user supplied SM count + int sm_count = hw_info.sm_count; + if (sm_count <= 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 = cutlass::round_up( + ceil_div(size<2>(problem_size), size<0>(tile_shape)), + size<0>(cluster_shape)); + int num_blocks = + num_m_blocks * size<0>(problem_size) * size<1>(problem_size); + + return Params{ + num_blocks, + {num_m_blocks}, + {size<0>(problem_size)}, + {size<1>(problem_size)}, + 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, bidh; + params.divmod_m_block(block_decode, m_block, block_decode); + params.divmod_b(block_decode, bidb, block_decode); + params.divmod_h(block_decode, bidh, block_decode); + return make_coord(m_block, _0{}, make_coord(bidb, bidh)); + } + + CUTLASS_DEVICE + PersistentTileScheduler& operator++() { + block_idx += gridDim.x; + return *this; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +template +struct TileSchedulerBwdAdapter { + using Params = typename Base::Params; + + Base base_; + + CUTLASS_DEVICE + TileSchedulerBwdAdapter(Params const& params) : base_(params) {} + + template + static Params to_underlying_arguments( + ProblemSize const& problem_size, + KernelHardwareInfo hw_info, + ClusterShape const& cluster_shape, + TileShape const& tile_shape) { + using namespace cute; + return Base::to_underlying_arguments( + select<0, 1, 3, 2, 4>(problem_size), + hw_info, + select<1, 0, 2>(cluster_shape), + select<1, 0, 2>(tile_shape)); + } + + static dim3 get_grid_shape(Params const& params) { + return Base::get_grid_shape(params); + } + + CUTLASS_DEVICE + bool is_valid() { + return base_.is_valid(); + } + + CUTLASS_DEVICE + auto get_block_coord() { + using namespace cute; + return select<1, 0, 2>(base_.get_block_coord()); + } + + CUTLASS_DEVICE + TileSchedulerBwdAdapter& operator++() { + ++base_; + return *this; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::fmha::kernel diff --git a/natten/csrc/include/natten/cuda/fna/epilogue/epilogue_pipelined.h b/natten/csrc/include/natten/cuda/fna/epilogue/epilogue_pipelined.h new file mode 100644 index 00000000..3b06c109 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna/epilogue/epilogue_pipelined.h @@ -0,0 +1,638 @@ +/* + * Copied from xFormers (https://github.com/facebookresearch/xformers/) + * + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + */ +/*************************************************************************************************** + * Copyright (c) 2017 - 2026 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 Epilogue for threadblock scoped GEMMs using Tensor Ops. + + File copied from "cutlass/epilogue/threadblock/epilogue.h" + then modified to: + (1) load 2 source fragments at the same time (pipelining) + (2) support reading from a different dtype + (3) pass the row id to the OutputOp if it takes it + (see MemoryEfficientAttentionNormalize) + Note that in general the fragment passed to the OutputOp could + span multiple rows but it does not happen with the configurations we have +*/ + +#pragma once + +#if defined(__CUDACC_RTC__) +#include +#else +#include +#endif + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include + +#include +#include +#include + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace epilogue { +namespace threadblock { + +template +struct ApplyEpilogueOp { + static CUTLASS_DEVICE typename Op::FragmentOutput apply( + Op const& output_op, + int row_id, + typename Op::FragmentAccumulator const& accum, + typename Op::FragmentOutput const& source) { + return output_op(accum, source); + } + static CUTLASS_DEVICE typename Op::FragmentOutput apply( + Op const& output_op, + int row_id, + typename Op::FragmentAccumulator const& accum) { + return output_op(accum); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Epilogue operator +template < + typename Shape_, ///< Shape of threadblock tile (concept: GemmShape) + typename WarpMmaOperator_, ///< Warp-level MMA operator (concept: + ///< gemm::warp::MmaTensorOp) + int PartitionsK, ///< Number of partitions of the K dimension + typename OutputTileIterator_, ///< Tile iterator writing output tensors + typename AccumulatorFragmentIterator_, ///< Fragment iterator selecting + ///< accumulators + typename WarpTileIterator_, ///< Warp-scoped tile iterator writing + ///< accumulators to SMEM + typename SharedLoadIterator_, ///< Threadblock-scoped tile iterator loading + ///< from SMEM + typename OutputOp_, ///< Output operator + typename Padding_, ///< Padding added to SMEM allocation to avoid bank + ///< conflicts (concept: MatrixShape) + int FragmentsPerPartition = + 1, ///< Used to coarsten the epilogue granularity + int IterationsUnroll = ///< Used to reduce binary size when epilogue op is + ///< large + (!IsEpilogueFunctorHeavy::value), + typename OutputTileSourceIterator_ = + OutputTileIterator_ ///< Tile iterator reading tensors + > +class EpiloguePipelined : public EpilogueBase< + Shape_, + typename WarpMmaOperator_::Shape, + PartitionsK, + AccumulatorFragmentIterator_, + WarpTileIterator_, + Padding_, + FragmentsPerPartition> { + public: + using Base = EpilogueBase< + Shape_, + typename WarpMmaOperator_::Shape, + PartitionsK, + AccumulatorFragmentIterator_, + WarpTileIterator_, + Padding_, + FragmentsPerPartition>; + + using Shape = Shape_; + using WarpMmaOperator = WarpMmaOperator_; + static int const kPartitionsK = PartitionsK; + using OutputTileIterator = OutputTileIterator_; + using OutputTileSourceIterator = OutputTileSourceIterator_; + using AccumulatorFragmentIterator = AccumulatorFragmentIterator_; + using WarpTileIterator = WarpTileIterator_; + using SharedLoadIterator = SharedLoadIterator_; + using OutputOp = OutputOp_; + using Padding = Padding_; + + using Layout = layout::RowMajor; + using LongIndex = typename Layout::LongIndex; + + /// The complete warp-level accumulator tile + using AccumulatorTile = typename Base::AccumulatorTile; + + /// Accumulator element + using ElementAccumulator = typename WarpTileIterator::Element; + + /// Output element + using ElementOutput = typename OutputTileIterator::Element; + using ElementSource = typename OutputTileSourceIterator::Element; + + /// Output access size + static int const kElementsPerAccess = OutputTileIterator::kElementsPerAccess; + + /// Tensor reference to destination tensor + using TensorRef = typename OutputTileIterator::TensorRef; + + /// Tensor reference to sync tensor + using SyncTensorRef = + typename cutlass::TensorRef; + + /// Const tensor reference to source tensor + using ConstTensorRef = typename OutputTileIterator::ConstTensorRef; + + /// Array type used to output + using OutputAccessType = Array< + typename OutputTileIterator::Element, + OutputTileIterator::kElementsPerAccess>; + using SourceAccessType = Array< + typename OutputTileSourceIterator::Element, + OutputTileSourceIterator::kElementsPerAccess>; + + /// Array type used by output functor + using AccumulatorAccessType = Array< + typename WarpTileIterator::Element, + OutputTileIterator::kElementsPerAccess>; + + /// Number of warps + using WarpCount = typename Base::WarpCount; + + static int constexpr kSmemTiles = Base::kFragmentsPerIteration > 1 + ? Base::kFragmentsPerIteration + : kPartitionsK; + static int constexpr kSmemPointerOffset = + Base::SharedStorage::StorageShape::kCount / kSmemTiles; + + public: + static_assert( + OutputTileSourceIterator::Fragment::kElements == + OutputTileIterator::Fragment::kElements, + "Mismatch between input tile and output tile iterator (kElements)"); + static_assert( + OutputTileSourceIterator::kIterations == OutputTileIterator::kIterations, + "Mismatch between input tile and output tile iterator (kIterations)"); + static_assert( + SharedLoadIterator::Fragment::kElements == + OutputTileIterator::Fragment::kElements, + "Mismatch between shared load iterator and output tile iterator."); + + static_assert( + OutputTileIterator::kElementsPerAccess, + "OutputTileIterator::kElementsPerAccess must not be zero."); + + static_assert( + !(OutputTileIterator::Fragment::kElements % + OutputTileIterator::kElementsPerAccess), + "Divisibility"); + + private: + /// Loads fragment from shared memory aligned with output tensor + SharedLoadIterator shared_load_iterator_; + + public: + /// Constructor + CUTLASS_DEVICE + EpiloguePipelined( + typename Base::SharedStorage& shared_storage, ///< Shared storage object + int thread_idx, ///< ID of a thread within the threadblock + int warp_idx, ///< ID of warp within threadblock + int lane_idx ///< Id of thread within warp + ) + : Base(shared_storage, thread_idx, warp_idx, lane_idx), + shared_load_iterator_(shared_storage.reference(), thread_idx) {} + + /// Streams the result to global memory + CUTLASS_DEVICE + void operator()( + OutputOp const& output_op, ///< Output operator + OutputTileIterator + destination_iterator, ///< Tile iterator for destination + AccumulatorTile const& + accumulators, ///< Complete warp-level accumulator tile + OutputTileSourceIterator + source_iterator) { ///< Threadblock tile coordinate in GEMM (in units + ///< of threadblock tiles) + + if (!output_op.is_source_needed()) { + compute_source_not_needed_(output_op, destination_iterator, accumulators); + } else { + compute_source_needed_( + output_op, destination_iterator, accumulators, source_iterator); + } + } + CUTLASS_DEVICE + void operator()( + OutputOp const& output_op, ///< Output operator + OutputTileIterator + destination_iterator, ///< Tile iterator for destination + AccumulatorTile const& + accumulators) { ///< Complete warp-level accumulator tile + compute_source_not_needed_(output_op, destination_iterator, accumulators); + } + + private: + template + struct acc2smem_source_not_needed; + + template + struct acc2smem_source_not_needed> { + template + CUTLASS_DEVICE static void helper( + AccumulatorFragmentIterator accum_fragment_iterator, + WarpTileIterator& warp_tile_iterator) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < Advance; i++) { + ++accum_fragment_iterator; + } + + CUTLASS_PRAGMA_UNROLL + for (int p = 0; p < Base::kFragmentsPerIteration; ++p) { + typename AccumulatorFragmentIterator::Fragment accum_fragment; + + accum_fragment_iterator.load(accum_fragment); + ++accum_fragment_iterator; + + warp_tile_iterator.store(accum_fragment); + if (p < Base::kFragmentsPerIteration - 1) { + warp_tile_iterator.add_pointer_offset(kSmemPointerOffset); + } + } + + if (Base::kFragmentsPerIteration > 1) { + warp_tile_iterator.add_pointer_offset( + kSmemPointerOffset * (1 - Base::kFragmentsPerIteration)); + } + } + + CUTLASS_DEVICE + static void push( + size_t pos, + AccumulatorFragmentIterator const& iterator_begin, + WarpTileIterator& warp_tile_iterator) { + int dummy[] = { + (pos == (Seq * Base::kFragmentsPerIteration)) && + (helper( + iterator_begin, warp_tile_iterator), + 0)...}; + + CUTLASS_UNUSED(dummy[0]); + } + }; + + static_assert( + kPartitionsK == 1 || Base::kFragmentsPerIteration == 1, + "One of these must be exactly 1."); + + /// Streams the result to global memory + CUTLASS_DEVICE + void compute_source_not_needed_( + OutputOp const& output_op, ///< Output operator + OutputTileIterator + destination_iterator, ///< Tile iterator for destination + AccumulatorTile const& + accumulators ///< Complete warp-level accumulator tile + ) { + // + // Iterator over warp-level accumulator fragment + // + + AccumulatorFragmentIterator accum_fragment_iterator(accumulators); + + // + // Iterate over accumulator tile + // + +#pragma unroll( \ + IterationsUnroll \ + ? OutputTileIterator::kIterations / Base::kFragmentsPerIteration \ + : 1) + for (int iter = 0; iter < OutputTileIterator::kIterations; + iter += Base::kFragmentsPerIteration) { + // + // Convert and store fragment + // + + __syncthreads(); + + acc2smem_source_not_needed>:: + push(iter, accum_fragment_iterator, this->warp_tile_iterator_); + + __syncthreads(); + + // + // Load fragments from shared memory + // + + CUTLASS_PRAGMA_UNROLL + for (int p = 0; p < Base::kFragmentsPerIteration; ++p) { + typename SharedLoadIterator::Fragment + aligned_accum_fragment[kPartitionsK]; + + shared_load_iterator_.load(aligned_accum_fragment[0]); + + if (p < Base::kFragmentsPerIteration - 1) { + shared_load_iterator_.add_pointer_offset(kSmemPointerOffset); + } else if (kPartitionsK > 1) { + plus add_fragments; + + CUTLASS_PRAGMA_UNROLL + for (int i = 1; i < kPartitionsK; ++i) { + shared_load_iterator_.add_pointer_offset(kSmemPointerOffset); + shared_load_iterator_.load(aligned_accum_fragment[i]); + aligned_accum_fragment[0] = add_fragments( + aligned_accum_fragment[0], aligned_accum_fragment[i]); + } + + shared_load_iterator_.add_pointer_offset( + (1 - kPartitionsK) * kSmemPointerOffset); + } + + // + // Compute the output result + // + + typename OutputTileIterator::Fragment output_fragment; + + apply_output_operator_source_not_needed_( + destination_iterator.thread_start_row(), + output_fragment, + output_op, + aligned_accum_fragment[0]); + + // + // Store the final result + // + + destination_iterator.store(output_fragment); + ++destination_iterator; + } + + if (Base::kFragmentsPerIteration > 1) { + shared_load_iterator_.add_pointer_offset( + kSmemPointerOffset * (1 - Base::kFragmentsPerIteration)); + } + } + } + + template + struct acc2smem_source_needed; + + template + struct acc2smem_source_needed> { + template + CUTLASS_DEVICE static void helper( + AccumulatorFragmentIterator accum_fragment_iterator, + WarpTileIterator& warp_tile_iterator) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < Advance; i++) { + ++accum_fragment_iterator; + } + + typename AccumulatorFragmentIterator::Fragment accum_fragment; + accum_fragment_iterator.load(accum_fragment); + warp_tile_iterator.store(accum_fragment); + } + + CUTLASS_DEVICE + static void push( + size_t pos, + AccumulatorFragmentIterator const& iterator_begin, + WarpTileIterator& warp_tile_iterator) { + int dummy[] = { + (pos == Seq) && + (helper(iterator_begin, warp_tile_iterator), 0)...}; + } + }; + + /// Streams the result to global memory + CUTLASS_DEVICE + void compute_source_needed_( + OutputOp const& output_op, ///< Output operator + OutputTileIterator + destination_iterator, ///< Tile iterator for destination + AccumulatorTile const& + accumulators, ///< Complete warp-level accumulator tile + OutputTileSourceIterator + source_iterator ///< Threadblock tile coordinate in GEMM (in units of + ///< threadblock tiles) + ) { + typename OutputTileSourceIterator::Fragment source_fragment[2]; + + source_fragment[0].clear(); + source_iterator.load(source_fragment[0]); + ++source_iterator; + source_fragment[1].clear(); + + // + // Iterator over warp-level accumulator fragment + // + + AccumulatorFragmentIterator accum_fragment_iterator(accumulators); + + // + // Iterate over accumulator tile + // + +#pragma unroll(IterationsUnroll ? OutputTileIterator::kIterations : 1) + for (int iter = 0; iter < OutputTileIterator::kIterations; ++iter) { + if (iter > 0) { + __syncthreads(); + } + // + // Load the source for next iteration (pipelining) + // + + if (iter + 1 < OutputTileIterator::kIterations) { + source_iterator.load(source_fragment[(iter + 1) % 2]); + } + ++source_iterator; + acc2smem_source_needed< + cutlass::make_index_sequence>:: + push(iter, accum_fragment_iterator, this->warp_tile_iterator_); + + __syncthreads(); + + // + // Load fragments from shared memory + // + + typename SharedLoadIterator::Fragment + aligned_accum_fragment[kPartitionsK]; + + shared_load_iterator_.load(aligned_accum_fragment[0]); + + // If the number of k-slices is > 1 - perform a reduction amongst the + // k-slices + if (kPartitionsK > 1) { + plus add_fragments; + + CUTLASS_PRAGMA_UNROLL + for (int i = 1; i < kPartitionsK; ++i) { + shared_load_iterator_.add_pointer_offset(kSmemPointerOffset); + shared_load_iterator_.load(aligned_accum_fragment[i]); + aligned_accum_fragment[0] = add_fragments( + aligned_accum_fragment[0], aligned_accum_fragment[i]); + } + + shared_load_iterator_.add_pointer_offset( + (1 - kPartitionsK) * kSmemPointerOffset); + } + + // + // Compute the output result + // + + typename OutputTileIterator::Fragment output_fragment; + + apply_output_operator_( + destination_iterator.thread_start_row(), + output_fragment, + output_op, + aligned_accum_fragment[0], + source_fragment[iter % 2]); + + // + // Store the final result + // + + destination_iterator.store(output_fragment); + ++destination_iterator; + } + } + + /// Helper to invoke the output functor over each vector of output + CUTLASS_DEVICE + void apply_output_operator_( + int begin_row, + typename OutputTileIterator::Fragment& output_fragment, + OutputOp const& output_op, ///< Output operator + typename SharedLoadIterator::Fragment const& aligned_accum_fragment, + typename OutputTileSourceIterator::Fragment const& source_fragment) { + OutputAccessType* output_frag_ptr = + reinterpret_cast(&output_fragment); + + AccumulatorAccessType const* compute_frag_ptr = + reinterpret_cast(&aligned_accum_fragment); + + SourceAccessType const* source_frag_ptr = + reinterpret_cast(&source_fragment); + + int const kOutputOpIterations = OutputTileIterator::Fragment::kElements / + OutputTileIterator::kElementsPerAccess; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kOutputOpIterations; ++i) { + // Call the output operator + output_frag_ptr[i] = ApplyEpilogueOp::apply( + output_op, + begin_row + getRowOffset(i * OutputTileIterator::kElementsPerAccess), + compute_frag_ptr[i], + source_frag_ptr[i]); + } + } + + /// Helper to invoke the output functor over each vector of output + CUTLASS_DEVICE + void apply_output_operator_source_not_needed_( + int begin_row, + typename OutputTileIterator::Fragment& output_fragment, + OutputOp const& output_op, ///< Output operator + typename SharedLoadIterator::Fragment const& aligned_accum_fragment) { + OutputAccessType* output_frag_ptr = + reinterpret_cast(&output_fragment); + + AccumulatorAccessType const* compute_frag_ptr = + reinterpret_cast(&aligned_accum_fragment); + + int const kOutputOpIterations = OutputTileIterator::Fragment::kElements / + OutputTileIterator::kElementsPerAccess; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kOutputOpIterations; ++i) { + // Call the output operator + output_frag_ptr[i] = ApplyEpilogueOp::apply( + output_op, + begin_row + getRowOffset(i * OutputTileIterator::kElementsPerAccess), + compute_frag_ptr[i]); + } + } + + // This should be constexpr, but it's only supported on c++14 + static int CUTLASS_HOST_DEVICE getRowOffset(int i) { + using ThreadMap = typename OutputTileIterator::ThreadMap; + + CUTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; + ++cluster) { + CUTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + CUTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + int row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + int frag_row_idx = + (row + + ThreadMap::Iterations::kRow * + (group + ThreadMap::Iterations::kGroup * cluster)); + CUTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; + ++column) { + int frag_idx = ThreadMap::kElementsPerAccess * + (frag_row_idx * ThreadMap::Iterations::kColumn + column); + if (i < frag_idx + ThreadMap::kElementsPerAccess) { + return row_offset; + } + } + } + } + } + return -1; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fna/epilogue/epilogue_rescale_output.h b/natten/csrc/include/natten/cuda/fna/epilogue/epilogue_rescale_output.h new file mode 100644 index 00000000..0499111d --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna/epilogue/epilogue_rescale_output.h @@ -0,0 +1,267 @@ +/* + * Copied from xFormers (https://github.com/facebookresearch/xformers/) + * + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * BSD 3-Clause License + * + * 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 names of Facebook, Deepmind Technologies, NYU, NEC + * Laboratories America and IDIAP Research Institute 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 OWNER 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 Epilogue for threadblock scoped GEMMs using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory + to match canonical tensor layouts in global memory. Epilogues support + conversion and reduction operations. + + This is a copy of cutlass/epilogue/threadblock/epilogue.h that can + handle "row_id" as a first argument, as uses it to get the corresponding + `m_prime` / `s_prime` to rescale the output. +*/ + +#pragma once + +#if defined(__CUDACC_RTC__) +#include +#else +#include +#endif + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include + +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace epilogue { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Applies a linear combination operator to an array of elements. +// output <- alpha * accumulator + beta * source +// with: +// alpha = 1 / s_prime (to normalize when isLast=True, 1 otherwise) +// beta = alpha / m_prime (renormalize the output when the max changes) +// source is the current output +template < + typename ElementOutput_, ///< Data type used to store tensors + typename ElementSource_, //< Data type for source (usually matches + //`ElementOutput`) + int Count, ///< Number of elements computed per operation. + ///< Usually it is 128/sizeof_bits, + ///< but we use 64 or 32 sometimes when there are not enough data + ///< to store + typename ElementAccumulator_, ///< Accumulator data type + typename ElementCompute_, ///< Data type used to compute linear combination + bool isFirst, + bool isLast, + typename FragmentAlphaBeta_, + FloatRoundStyle Round = FloatRoundStyle::round_to_nearest> +class MemoryEfficientAttentionNormalize { + public: + using ElementOutput = ElementOutput_; + using ElementSource = ElementSource_; + using ElementAccumulator = ElementAccumulator_; + using ElementCompute = ElementCompute_; + + static int const kCount = Count; + + using FragmentOutput = Array; + using FragmentSource = Array; + using FragmentAccumulator = Array; + using ComputeFragment = Array; + using FragmentAlphaBeta = FragmentAlphaBeta_; + + static FloatRoundStyle const kRound = Round; + + private: + // + // Data members + // + + FragmentAlphaBeta const& s_prime_; + FragmentAlphaBeta const& m_prime_; + + public: + /// Constructs the function object, possibly loading from pointers in host + /// memory + CUTLASS_HOST_DEVICE + MemoryEfficientAttentionNormalize( + FragmentAlphaBeta const& s_prime, + FragmentAlphaBeta const& m_prime) + : s_prime_(s_prime), m_prime_(m_prime) {} + + /// Returns true if source is needed + CUTLASS_HOST_DEVICE + bool is_source_needed() const { + return !isFirst; + } + + /// Functionally required for serial reduction in the epilogue + CUTLASS_HOST_DEVICE + void set_k_partition(int k_partition, int k_partition_count) {} + + /// Computes linear scaling: D = alpha * accumulator + beta * source + CUTLASS_HOST_DEVICE + FragmentOutput operator()( + int row, + FragmentAccumulator const& accumulator, + FragmentSource const& source) const { + assert(!isFirst); + + // Convert source to interal compute numeric type + NumericArrayConverter + source_converter; + NumericArrayConverter + accumulator_converter; + + // Convert to destination numeric type + NumericArrayConverter + destination_converter; + + ComputeFragment converted_source = source_converter(source); + ComputeFragment converted_accumulator = accumulator_converter(accumulator); + + // Perform binary operations + ComputeFragment intermediate; + + multiplies mul_add_source; + multiply_add mul_add_accumulator; + + ElementCompute alpha = isLast ? (1 / s_prime_[row]) : 1; + ElementCompute beta = alpha * m_prime_[row]; + + intermediate = mul_add_source(beta, converted_source); // X = beta * C + + intermediate = mul_add_accumulator( + alpha, converted_accumulator, intermediate); // D = alpha * Accum + X + + return destination_converter(intermediate); + } + + /// Computes linear scaling: D = alpha * accumulator + CUTLASS_HOST_DEVICE + FragmentOutput operator()(int row, FragmentAccumulator const& accumulator) + const { + assert(isFirst); + + // Convert source to interal compute numeric type + NumericArrayConverter + accumulator_converter; + + // Convert to destination numeric type + NumericArrayConverter + destination_converter; + + ComputeFragment converted_accumulator = accumulator_converter(accumulator); + + ComputeFragment intermediate; + multiplies mul_accumulator; + + ElementCompute alpha = isLast ? (1 / s_prime_[row]) : 1; + + intermediate = mul_accumulator( + alpha, converted_accumulator); // X = alpha * C + uniform + + return destination_converter(intermediate); + } +}; + +} // namespace thread + +namespace threadblock { +template < + typename EO, + typename ES, + int Count, + typename EA, + typename EC, + bool F, + bool L, + typename FAB, + FloatRoundStyle R> +struct ApplyEpilogueOp> { + using Op = thread:: + MemoryEfficientAttentionNormalize; + static CUTLASS_DEVICE typename Op::FragmentOutput apply( + Op const& output_op, + int row_id, + typename Op::FragmentAccumulator const& accum, + typename Op::FragmentSource const& source) { + return output_op(row_id, accum, source); + } + static CUTLASS_DEVICE typename Op::FragmentOutput apply( + Op const& output_op, + int row_id, + typename Op::FragmentAccumulator const& accum) { + return output_op(row_id, accum); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fna/epilogue/epilogue_thread_apply_logsumexp.h b/natten/csrc/include/natten/cuda/fna/epilogue/epilogue_thread_apply_logsumexp.h new file mode 100644 index 00000000..8e1b72dd --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna/epilogue/epilogue_thread_apply_logsumexp.h @@ -0,0 +1,184 @@ +/* + * Copied from xFormers (https://github.com/facebookresearch/xformers/) + * + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + */ +/*************************************************************************************************** + * Copyright (c) 2017 - 2026 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 Functor performing linear combination operations used by epilogues. +*/ + +#pragma once + +#include + +#include "cutlass/array.h" +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/thread/activation.h" +#include "cutlass/functional.h" +#include "cutlass/numeric_conversion.h" +#include "cutlass/numeric_types.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace epilogue { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace detail { + +template +struct ArrayExponential { + CUTLASS_HOST_DEVICE + Array operator()( + Array const& input) const { + Array result; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < ElementsPerAccess; ++i) { + result[i] = expf(input[i]); + } + + return result; + } +}; + +template +struct ArrayExponential { + CUTLASS_DEVICE + Array operator()( + Array const& input) const { + Array result; + + int const kVectorCount = ElementsPerAccess / 2; + + __half2 const* input_ptr = + reinterpret_cast<__half2 const*>(input.raw_data()); + __half2* res_ptr = reinterpret_cast<__half2*>(result.raw_data()); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kVectorCount; ++i) { + res_ptr[i] = h2exp(input_ptr[i]); + } + + return result; + } +}; +} // namespace detail + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Applies: +/// output <- (input).exp() +template < + typename ElementOutput_, // output + typename ElementLSE_, // accumulator from LSE + typename ElementAccumulator_, // accumulator from matmul + typename ElementCompute_, // intermediate compute (and exp calculation) + int ElementsPerAccess> +// class ApplyLogSumExp { +class ApplyExp { + public: + using ElementOutput = ElementOutput_; + using ElementAccumulator = ElementAccumulator_; + using ElementCompute = ElementCompute_; + using ElementLSE = ElementLSE_; + + static int const kElementsPerAccess = ElementsPerAccess; + static int const kCount = kElementsPerAccess; + static const ScaleType::Kind kScale = + cutlass::epilogue::thread::ScaleType::NoBetaScaling; + + using FragmentOutput = Array; + using FragmentAccumulator = Array; + using FragmentCompute = Array; + using FragmentLSE = Array; + using FragmentScaleBias = FragmentLSE; // Used by epilogue_smem_accumulator.h + + public: + // + // Methods + // + + CUTLASS_HOST_DEVICE + ApplyExp() {} + + /// Returns true if source is needed + CUTLASS_HOST_DEVICE + bool is_source_needed() const { + return true; + } + + /// Functionally required for serial reduction in the epilogue + CUTLASS_HOST_DEVICE + void set_k_partition(int k_partition, int k_partition_count) {} + + CUTLASS_HOST_DEVICE + FragmentOutput operator()( + FragmentAccumulator const& AB/*, + FragmentLSE const& scale_unused, + // bias used as LSE + FragmentLSE const& bias*/) const { + FragmentCompute frag_AB = NumericArrayConverter< + ElementCompute, + ElementAccumulator, + kElementsPerAccess>()(AB); + // FragmentCompute frag_lse_compute = + // NumericArrayConverter()( + // bias); + FragmentCompute frag_compute; + + // minus minus_lse; + detail::ArrayExponential apply_exp; + // frag_compute = minus_lse(frag_AB, frag_lse_compute); + // frag_compute = apply_exp(frag_compute); + frag_compute = apply_exp(frag_AB); + + return NumericArrayConverter< + ElementOutput, + ElementCompute, + kElementsPerAccess>()(frag_compute); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace epilogue +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fna/epilogue/predicated_tile_iterator.h b/natten/csrc/include/natten/cuda/fna/epilogue/predicated_tile_iterator.h new file mode 100644 index 00000000..18c47b8b --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna/epilogue/predicated_tile_iterator.h @@ -0,0 +1,628 @@ +/* + * Copied from CUTLASS (https://github.com/NVIDIA/cutlass/) and edited. + */ +/*************************************************************************************************** + * Copyright (c) 2017 - 2026 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 Epilogue for threadblock scoped GEMMs using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory + to match canonical tensor layouts in global memory. Epilogues support + conversion and reduction operations. + +*/ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { + +//////////////////////////////////////////////////////////////////////////////// + +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator used to load and store output tile from global memory in +/// epilogue. +/// +/// Satisfies: ReadableTileIterator | CustomPredicatedTileIterator | +/// ForwardTileIterator +/// +template < + int NADim, + typename ThreadMap_, ///< Thread map (conept: OutputTileThreadMap) + typename Element_ ///< Element data type + > +class CustomPredicatedTileIterator { + public: + static_assert(NADim >= 1 && NADim < 4); + using Dim = typename natten::cuda::fna::GetDim::type; + + using ThreadMap = ThreadMap_; + using Shape = typename ThreadMap::Shape; + + using Element = Element_; + + using Layout = layout::RowMajor; + using TensorRef = TensorRef; + using ConstTensorRef = typename TensorRef::ConstTensorRef; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using TensorCoord = MatrixCoord; + + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + static int const kThreads = ThreadMap::kThreads; + static int const kIterations = ThreadMap::Count::kTile; + + static_assert( + ThreadMap::Iterations::kRow > 0, + "ThreadMap::Iterations::kRow must be > 0"); + static_assert( + ThreadMap::Iterations::kGroup > 0, + "ThreadMap::Iterations::kGroup must be > 0"); + static_assert( + ThreadMap::Iterations::kCluster > 0, + "ThreadMap::Iterations::kCluster must be > 0"); + static_assert( + ThreadMap::Iterations::kColumn > 0, + "ThreadMap::Iterations::kColumn must be > 0"); + + /// Fragment object + using Fragment = Array< + Element, + ThreadMap::Iterations::kColumn * ThreadMap::Iterations::kRow * + ThreadMap::Iterations::kGroup * ThreadMap::Iterations::kCluster * + ThreadMap::kElementsPerAccess>; + + /// Memory access size + using AccessType = AlignedArray; + + // + // Parameters struct + // + using ParamsBase = CustomPredicatedTileIteratorParams; + + /// Uses a non-template class + struct Params : ParamsBase { + using Base = ParamsBase; + + CUTLASS_HOST_DEVICE + Params() {} + + CUTLASS_HOST_DEVICE + Params(Dim stride, Dim extent_row) + : ParamsBase( + stride, + int32_t(sizeof(AccessType)) / kElementsPerAccess, + extent_row, + make_OutputTileThreadMapDesc()) {} + + CUTLASS_HOST_DEVICE + Params(Base const& base) : Base(base) {} + }; + + /// Mask object + struct Mask { + static int const kCount = ThreadMap::Iterations::kColumn; + + /// Predicate state + bool predicates[kCount]; + + // + // Mask + // + CUTLASS_HOST_DEVICE + Mask() { + enable(); + } + + ///< Efficiently disables all accesses guarded by mask + CUTLASS_HOST_DEVICE void clear() { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + predicates[i] = false; + } + } + + ///< CUTLASS_HOST_DEVICE enables all accesses guarded by mask + CUTLASS_DEVICE void enable() { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + predicates[i] = true; + } + } + }; + + private: + // + // Data members + // + + /// Parameters structure containing reference and precomputed state. + ParamsBase params_; + + /// Byte-level pointer. This pointer is usually for both load() and store(), + /// unless PermuteD is performed. When having PermuteD, byte_pointer_ is only + /// for load(). + uint8_t* byte_pointer_; + + // /// Byte-level pointer for store(). Due to PermuteD Op, store_byte_pointer_ + // may be with different address computation compared to byte_pointer_. + // uint8_t *store_byte_pointer_; + + /// Array of boolean values to contain steady-state predicates + Mask mask_; + + /// Extent of the matrix tile in rows + Dim extent_row_; + int32_t extent_row_int; + + /// Extent of the matrix tile in rows + Index extent_column_; + + /// A thread's starting row position (assuming steady-state predicates have + /// been computed) + Index thread_start_row_; + + /// A thread's starting column + Index thread_start_column_; + + /// Internal state counter + int state_[3]; + + // + // Static asserts about internal strides + // + + static_assert(sizeof(extent_column_) == 4, "Expected 32b extents"); + static_assert(sizeof(thread_start_row_) == 4, "Expected 32b extents"); + + private: + // + // Methods + // + + public: + // + // Methods + // + + /// Constructor + CUTLASS_DEVICE + CustomPredicatedTileIterator( + ParamsBase const& params, + Element* pointer, + Dim extent_row, + int32_t extent_col, + int thread_idx, + TensorCoord threadblock_offset = TensorCoord()) + : params_(params) { + TensorCoord thread_offset = + ThreadMap::initial_offset(thread_idx) + threadblock_offset; + + extent_row_ = extent_row; + extent_column_ = extent_col; + extent_row_int = extent_row.prod32(); + + thread_start_row_ = thread_offset.row(); + thread_start_column_ = thread_offset.column(); + + // Initialize predicates + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kColumn; ++c) { + mask_.predicates[c] = + ((thread_offset.column() + ThreadMap::Delta::kColumn * c) < + extent_col); + } + + // Null pointer performs no accesses + if (!pointer) { + mask_.clear(); + } + + // if (ScatterD && !indices) { + // mask_.clear(); + // } + + // Initialize byte_pointer_ + byte_pointer_ = reinterpret_cast(pointer) + + //(natten::cuda::fna::map_index_to_coord(thread_offset.row(), + // extent_row_) * params_.stride).sum() + + LongIndex(thread_offset.column()) * sizeof(AccessType) / + kElementsPerAccess; + + // if (ScatterD) { + // byte_pointer_ = reinterpret_cast(pointer) + + // LongIndex(thread_offset.column()) * sizeof(AccessType) / + // kElementsPerAccess; + //} + + // store_byte_pointer_ is set to be the same with byte_pointer_ unless + // PermuteD is used. + // store_byte_pointer_ = PermuteD ? reinterpret_cast(pointer) : + // byte_pointer_; + // store_byte_pointer_ = byte_pointer_; + + // Initialize internal state counter + state_[0] = state_[1] = state_[2] = 0; + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + // store_byte_pointer_ += pointer_offset * sizeof_bits::value / 8; + byte_pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_byte_offset(Fragment& frag, int64_t byte_offset) const { + uint8_t* byte_pointer = byte_pointer_; + AccessType* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; + ++cluster) { + CUTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + CUTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + int frag_row_idx = + (row + + ThreadMap::Iterations::kRow * + (group + ThreadMap::Iterations::kGroup * cluster)); + + int32_t row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + + auto row_offset_full_ = natten::cuda::fna::map_index_to_coord( + row_offset + thread_start_row_, extent_row_); + // bool row_guard = + // natten::cuda::fna::is_coord_within_upper_bound(row_offset_full_, + // extent_row_); + bool row_guard = row_offset + thread_start_row_ < extent_row_int; + + AccessType* memory_pointer = reinterpret_cast( + byte_pointer + byte_offset + + (row_offset_full_ * params_.stride).sum()); + + // if (ScatterD && row_guard) { + // assert(indices_); + + // memory_pointer = reinterpret_cast(byte_pointer + + // byte_offset + + // LongIndex(indices_[row_offset + thread_start_row_]) * + // LongIndex(params_.stride)); + //} + + CUTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; + ++column) { + bool guard = row_guard && mask_.predicates[column]; + + cutlass::arch::global_load( + frag_ptr + [frag_row_idx * ThreadMap::Iterations::kColumn + column], + (void*)&memory_pointer + [column * ThreadMap::Delta::kColumn / kElementsPerAccess], + guard); + } + + if (row + 1 < ThreadMap::Iterations::kRow) { + // if (!ScatterD) { + // byte_pointer += params_.increment_row; + //} + } + } + + if (group + 1 < ThreadMap::Iterations::kGroup) { + // byte_pointer += params_.increment_group; + } + } + + if (cluster + 1 < ThreadMap::Iterations::kCluster) { + // byte_pointer += params_.increment_cluster; + } + } + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) const { + load_with_byte_offset(frag, 0); + } + + /// Stores a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, int64_t byte_offset) const { + // uint8_t *byte_pointer = store_byte_pointer_; + uint8_t* byte_pointer = byte_pointer_; + AccessType const* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; + ++cluster) { + CUTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + CUTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + int frag_row_idx = + (row + + ThreadMap::Iterations::kRow * + (group + ThreadMap::Iterations::kGroup * cluster)); + + int32_t row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + + auto row_offset_full_ = natten::cuda::fna::map_index_to_coord( + row_offset + thread_start_row_, extent_row_); + // bool row_guard = + // natten::cuda::fna::is_coord_within_upper_bound(row_offset_full_, + // extent_row_); + bool row_guard = row_offset + thread_start_row_ < extent_row_int; + + AccessType* memory_pointer = reinterpret_cast( + byte_pointer + byte_offset + + (row_offset_full_ * params_.stride).sum()); + + // if (ScatterD && row_guard) { + // assert(indices_); + + // memory_pointer = reinterpret_cast(byte_pointer + + // byte_offset + + // LongIndex(indices_[row_offset + thread_start_row_]) * + // LongIndex(params_.stride)); + //} + + CUTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; + ++column) { + bool guard = row_guard && mask_.predicates[column]; + + // if (PermuteD) { + + // int col_offset = column * ThreadMap::Delta::kColumn; + + // int col = col_offset + thread_start_column_; + // int row = row_offset + thread_start_row_; + + // // Locate memory_pointer + // memory_pointer = reinterpret_cast(byte_pointer + + // byte_offset + // + permute_layout_(PitchLinearCoord(col, row)) * + // sizeof(AccessType) / kElementsPerAccess); + // } + + // if (UseCUDAStore) { + // if (guard) { + // memory_pointer[0] = + // frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + + // column]; + // } + //} else { + cutlass::arch::global_store( + frag_ptr + [frag_row_idx * ThreadMap::Iterations::kColumn + column], + (void*)&memory_pointer[0], + guard); + //} + + // if (!PermuteD) { + memory_pointer += (ThreadMap::Delta::kColumn / kElementsPerAccess); + //} + } + + if (row + 1 < ThreadMap::Iterations::kRow) { + // if (!ScatterD && !PermuteD) { + // byte_pointer += params_.increment_row; + //} + } + } + + if (group + 1 < ThreadMap::Iterations::kGroup) { + // byte_pointer += params_.increment_group; + } + } + + if (cluster + 1 < ThreadMap::Iterations::kCluster) { + // byte_pointer += params_.increment_cluster; + } + } + } + + /// Stores a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) const { + store_with_byte_offset(frag, 0); + } + + CUTLASS_DEVICE + MatrixCoord thread_start() const { + return MatrixCoord(thread_start_row_, thread_start_column_); + } + + /// Need to get the thread start row from the tile iterator + CUTLASS_DEVICE + int32_t thread_start_row() const { + return thread_start_row_; + } + + /// Need to get the thread start row from the tile iterator + CUTLASS_DEVICE + int32_t thread_start_column() const { + return thread_start_column_; + } + + /// Extent of the matrix in columns + CUTLASS_DEVICE + Index extent_column() const { + return extent_column_; + } + + /// Advances to the next position to load or store + CUTLASS_HOST_DEVICE + CustomPredicatedTileIterator& operator++() { + ++state_[0]; + + // if (!ScatterD && !PermuteD) { + // store_byte_pointer_ += params_.advance_row; + //} + + // if (!ScatterD) { + // byte_pointer_ += params_.advance_row; + //} + + thread_start_row_ += ThreadMap::Shape::kRow; + + if (state_[0] == ThreadMap::Count::kRow) { + state_[0] = 0; + ++state_[1]; + // byte_pointer_ += params_.advance_group; + // store_byte_pointer_ += params_.advance_group; + + thread_start_row_ += (ThreadMap::Shape::kGroup - 1) * + ThreadMap::Shape::kRow * ThreadMap::Count::kRow; + + if (state_[1] == ThreadMap::Count::kGroup) { + state_[1] = 0; + ++state_[2]; + // byte_pointer_ += params_.advance_cluster; + // store_byte_pointer_ += params_.advance_cluster; + + thread_start_row_ += ThreadMap::Count::kGroup * + ThreadMap::Shape::kGroup * ThreadMap::Count::kRow * + ThreadMap::Shape::kRow; + + if (state_[2] == ThreadMap::Count::kCluster) { + state_[2] = 0; + // byte_pointer_ += params_.advance_tile; + // store_byte_pointer_ += params_.advance_tile; + + thread_start_row_ += ThreadMap::Shape::kGroup * + ThreadMap::Shape::kRow * ThreadMap::Shape::kCluster * + ThreadMap::Shape::kTile; + } + } + } + + return *this; + } + + /// Advances a number of positions to load or store + CUTLASS_HOST_DEVICE + CustomPredicatedTileIterator& operator+=(int increment) { + // Row + state_[0] += increment; + int increment_row = state_[0] / ThreadMap::Count::kRow; + state_[0] = state_[0] % ThreadMap::Count::kRow; + + // byte_pointer_ += (params_.advance_row * increment); + // store_byte_pointer_ += (params_.advance_row * increment); + thread_start_row_ += (ThreadMap::Shape::kRow * increment); + + // Group + state_[1] += increment_row; + int increment_group = state_[1] / ThreadMap::Count::kGroup; + state_[1] = state_[1] % ThreadMap::Count::kGroup; + + // byte_pointer_ += (params_.advance_group * increment_row); + // store_byte_pointer_ += (params_.advance_group * increment_row); + thread_start_row_ += (ThreadMap::Shape::kGroup - 1) * + ThreadMap::Shape::kRow * ThreadMap::Count::kRow * increment_row; + + // Cluster + state_[2] += increment_group; + int increment_cluster = state_[2] / ThreadMap::Count::kCluster; + state_[2] = state_[2] % ThreadMap::Count::kCluster; + + // byte_pointer_ += (params_.advance_cluster * increment_group); + // store_byte_pointer_ += (params_.advance_cluster * increment_group); + thread_start_row_ += ThreadMap::Count::kGroup * ThreadMap::Shape::kGroup * + ThreadMap::Count::kRow * ThreadMap::Shape::kRow * increment_group; + + // Tile + // byte_pointer_ += (params_.advance_tile * increment_cluster); + // store_byte_pointer_ += (params_.advance_tile * increment_cluster); + thread_start_row_ += ThreadMap::Shape::kGroup * ThreadMap::Shape::kRow * + ThreadMap::Shape::kCluster * ThreadMap::Shape::kTile * + increment_cluster; + + return *this; + } + + ///< Efficiently disables all accesses guarded by mask + CUTLASS_DEVICE void clear_mask() { + mask_.clear(); + } + + ///< Efficiently enables all accesses guarded by mask + CUTLASS_DEVICE void enable_mask() { + mask_.enable(); + } + + ///< Sets the mask + CUTLASS_DEVICE void get_mask(Mask& mask) const { + mask = mask_; + } + + ///< Sets the mask + CUTLASS_DEVICE void set_mask(Mask const& mask) { + mask_ = mask; + } +}; + +} // namespace threadblock +} // namespace epilogue +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fna/epilogue/predicated_tile_iterator_params.h b/natten/csrc/include/natten/cuda/fna/epilogue/predicated_tile_iterator_params.h new file mode 100644 index 00000000..1bdd3552 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna/epilogue/predicated_tile_iterator_params.h @@ -0,0 +1,152 @@ +/* + * Copied from CUTLASS (https://github.com/NVIDIA/cutlass/) and edited. + */ +/*************************************************************************************************** + * Copyright (c) 2017 - 2026 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 +*/ + +#pragma once + +#include + +#include +#include + +#include + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace epilogue { +namespace threadblock { + +// +// Parameters struct for PredicatedTileIterator +// + +template +struct CustomPredicatedTileIteratorParams { + using Dim = typename natten::cuda::fna::GetDim::type; + + using Index = int32_t; + using LongIndex = int64_t; + + // + // Data members + // + + Dim stride; ///< stride in bytes between rows + + // LongIndex increment_row; ///< increment quantity (in bytes) to + // advance when moving between rows LongIndex increment_group; ///< + // increment quantity (in bytes) to advance when moving to the next group + // LongIndex increment_cluster; ///< increment quantity (in bytes) to + // advance when moving to the next cluster + + // LongIndex advance_row; ///< amount to add to move to the next + // 'row' position LongIndex advance_group; ///< amount to add to move + // to the next 'group' position LongIndex advance_cluster; ///< amount to + // add to move to the next 'cluster' position LongIndex advance_tile; ///< + // amount to add to move to the next 'tile' + + // + // Methods + // + + CUTLASS_HOST_DEVICE + Status initialize( + Dim stride_, + Index elem, + Dim extent_row, + OutputTileThreadMapDesc thread_map) { + stride = stride_ * elem; + + // increment_row = (stride * + // natten::cuda::fna::map_index_to_coord(thread_map.delta.row, + // extent_row)).sum(); + + // increment_group = (stride * + // natten::cuda::fna::map_index_to_coord(thread_map.delta.group, extent_row) + // - stride * natten::cuda::fna::map_index_to_coord(thread_map.delta.row * + // (thread_map.iterations.row - 1), extent_row)).sum(); + + // increment_cluster = (stride * + // natten::cuda::fna::map_index_to_coord(thread_map.delta.cluster, + // extent_row) + // - stride * natten::cuda::fna::map_index_to_coord(thread_map.delta.group + // * (thread_map.iterations.group - 1), extent_row) + // - stride * natten::cuda::fna::map_index_to_coord(thread_map.delta.row * + // (thread_map.iterations.row - 1), extent_row)).sum(); + + // advance_row = (stride * + // natten::cuda::fna::map_index_to_coord(thread_map.shape.row, + // extent_row)).sum(); + + // advance_group = + // (stride * + // natten::cuda::fna::map_index_to_coord((thread_map.shape.group - 1) * + // thread_map.shape.row * thread_map.count.row, extent_row)).sum(); + // + // advance_cluster = ( + // stride * + // natten::cuda::fna::map_index_to_coord(thread_map.count.group * + // thread_map.shape.group * + // thread_map.count.row * + // thread_map.shape.row, extent_row)).sum(); + // + // advance_tile =( + // stride * + // natten::cuda::fna::map_index_to_coord(thread_map.shape.group * + // thread_map.shape.row * + // thread_map.shape.cluster * + // thread_map.shape.tile, extent_row)).sum(); + + return Status::kSuccess; + } + + CUTLASS_HOST_DEVICE + CustomPredicatedTileIteratorParams( + Dim stride, + Index elem, + Dim extent_row, + OutputTileThreadMapDesc thread_map) { + initialize(stride, elem, extent_row, thread_map); + } +}; + +} // namespace threadblock +} // namespace epilogue +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fna/fna_backward.cuh b/natten/csrc/include/natten/cuda/fna/fna_backward.cuh new file mode 100644 index 00000000..1cba0468 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna/fna_backward.cuh @@ -0,0 +1,219 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Fused Neighborhood Attention kernels are heavily based on the + * memory-efficient attention kernels from the xFormers project by Meta + * Platforms, Inc. + * + * Copyright (c) Facebook, Inc. and its affiliates + * + * BSD 3-Clause License + * + * 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 names of Facebook, Deepmind Technologies, NYU, NEC + * Laboratories America and IDIAP Research Institute 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 OWNER 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 + +#include + +#include + +namespace natten { +namespace cuda { +namespace fna { + +template < + typename IntTuple, + typename BooleanTuple, + typename MemoryAllocator // In lieu of a caching allocator, and while we + // only bind with torch + > +void fna_backward_generic( + at::ScalarType dtype, + const int cc, + const size_t max_smem, + cudaStream_t stream, + MemoryAllocator alloc_bytes, + void* grad_out_ptr, + void* query_ptr, + void* key_ptr, + void* value_ptr, + // bias is not supported! + void* logsumexp_ptr, + void* delta_ptr, + void* out_ptr, + // void* workspace_ptr, + // Outputs: + void* grad_query_ptr, + void* grad_key_ptr, + void* grad_value_ptr, + // Params + int32_t batch_size, + IntTuple spatial_extent, + int32_t heads, + int32_t dim, + int32_t dim_value, + IntTuple kernel_size, + IntTuple stride, + IntTuple dilation, + BooleanTuple is_causal, + float attn_scale, + IntTuple query_tile_shape, + IntTuple key_tile_shape, + IntTuple num_splits_key) { + static constexpr auto kRank = + std::tuple_size::value; + using Dim = typename GetDim::type; + + bool kernel_launched = false; + const auto maxK = std::max(dim, dim_value); + + auto launchKernel = [&](auto _k, auto kernel_fn) { + using Kernel = decltype(_k); + using scalar_t = typename Kernel::scalar_t; + (void)_k; + + if (kernel_launched) { + return; + } + // Check if this kernel is compatible + if (Kernel::kMaxK < maxK) { + return; + } + + if (natten::flatten(query_tile_shape) != Kernel::kBlockSizeI || + natten::flatten(key_tile_shape) != Kernel::kBlockSizeJ) { + return; + } + + if ((dim % Kernel::kMinimumAlignment) || + (dim % Kernel::kMinimumAlignment) || + (dim_value % Kernel::kMinimumAlignment)) { + return; + } + // Uses too much shmem + size_t smem_bytes = sizeof(typename Kernel::SharedStorage); + if (smem_bytes > max_smem) { + return; + } + + kernel_launched = true; + + typename Kernel::Params p; + p.query_ptr = (scalar_t*)query_ptr; + p.key_ptr = (scalar_t*)key_ptr; + p.value_ptr = (scalar_t*)value_ptr; + p.logsumexp_ptr = (typename Kernel::lse_scalar_t*)logsumexp_ptr; + p.output_ptr = (scalar_t*)out_ptr; + p.grad_output_ptr = (scalar_t*)grad_out_ptr; + p.grad_query_ptr = (scalar_t*)grad_query_ptr; + p.grad_key_ptr = (scalar_t*)grad_key_ptr; + p.grad_value_ptr = (scalar_t*)grad_value_ptr; + p.delta_ptr = (float*)delta_ptr; + + p.scale = attn_scale; + + p.num_heads = heads; + p.head_dim = dim; + p.head_dim_value = dim_value; + p.num_queries = tuple_to_na_dim(spatial_extent); + p.num_batches = batch_size; + + p.kernel_size = tuple_to_na_dim(kernel_size); + p.stride = tuple_to_na_dim(stride); + p.dilation = tuple_to_na_dim(dilation); + + p.query_tile_shape = tuple_to_na_dim(query_tile_shape); + p.key_tile_shape = tuple_to_na_dim(key_tile_shape); + + p.num_splits_key = tuple_to_na_dim(num_splits_key); + + int64_t size_bytes = p.workspace_size(); + if (size_bytes) { + void* workspace_ptr = nullptr; + alloc_bytes( + &workspace_ptr, size_bytes, true /*p.should_zero_workspace()*/); + p.workspace = (float*)workspace_ptr; + } + + Kernel::check_supported(p); + + if (smem_bytes > 0xc000) { + // https://docs.nvidia.com/cuda/cuda-c-programming-guide/#features-and-technical-specifications-technical-specifications-per-compute-capability + auto err = cudaFuncSetAttribute( + kernel_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes); + NATTEN_CHECK( + err != cudaErrorInvalidValue, + "This GPU does not have enough shared-memory."); + NATTEN_CUDA_CHECK(err); + } + + // second syntax resulted in the error below on windows + // error C3495: 'kernel_fn': a simple capture must be a variable + // with automatic storage duration declared + // in the reaching scope of the lambda +#ifdef NATTEN_WINDOWS + cudaFuncAttributes attr; + NATTEN_CUDA_CHECK(cudaFuncGetAttributes(&attr, kernel_fn)); + NATTEN_CHECK( + attr.binaryVersion >= Kernel::ArchTag::kMinComputeCapability, + "Something went wrong in the build process"); +#else + auto checkBinaryArchMatches = [&]() { + cudaFuncAttributes attr; + NATTEN_CUDA_CHECK(cudaFuncGetAttributes(&attr, kernel_fn)); + return attr.binaryVersion >= Kernel::ArchTag::kMinComputeCapability; + }; + NATTEN_CHECK( + checkBinaryArchMatches(), "Something went wrong in the build process"); +#endif + + kernel_fn<<>>(p); + + // if (size_bytes) { + // // NATTEN_CHECK( + // // accum_ptr != nullptr, "Expected accum_ptr to be set, got + // // nullptr."); + // NATTEN_CUDA_CHECK(cudaFreeAsync(workspace_ptr, stream)); + // } + }; + + DISPATCH_FNA_BACKWARD_KERNEL(kRank, cc, dtype, is_causal, launchKernel); + NATTEN_CHECK( + kernel_launched, + "Could not find a compatible fused neighborhood attention backward kernel."); + NATTEN_CUDA_CHECK(cudaGetLastError()); +} + +} // namespace fna +} // namespace cuda +} // namespace natten diff --git a/natten/csrc/include/natten/cuda/fna/fna_forward.cuh b/natten/csrc/include/natten/cuda/fna/fna_forward.cuh new file mode 100644 index 00000000..bf129ae8 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna/fna_forward.cuh @@ -0,0 +1,192 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Fused Neighborhood Attention kernels are heavily based on the + * memory-efficient attention kernels from the xFormers project by Meta + * Platforms, Inc. + * + * Copyright (c) Facebook, Inc. and its affiliates + * + * BSD 3-Clause License + * + * 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 names of Facebook, Deepmind Technologies, NYU, NEC + * Laboratories America and IDIAP Research Institute 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 OWNER 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 + +#include + +#include + +namespace natten { +namespace cuda { +namespace fna { + +template < + typename IntTuple, + typename BooleanTuple, + typename MemoryAllocator // In lieu of a caching allocator, and while we + // only bind with torch + > +void fna_forward_generic( + at::ScalarType dtype, + const int cc, + const size_t max_smem, + cudaStream_t stream, + MemoryAllocator alloc_bytes, + void* query_ptr, + void* key_ptr, + void* value_ptr, + void* out_ptr, + int32_t batch_size, + IntTuple spatial_extent, + int32_t heads, + int32_t dim, + int32_t dim_value, + IntTuple kernel_size, + IntTuple stride, + IntTuple dilation, + BooleanTuple is_causal, + float attn_scale, + void* logsumexp_ptr, + IntTuple query_tile_shape, + IntTuple key_tile_shape) { + static constexpr auto kRank = + std::tuple_size::value; + using Dim = typename GetDim::type; + + bool compute_logsumexp = logsumexp_ptr != nullptr; + bool kernel_launched = false; + auto launchKernel = [&](auto _k, auto kernel_fn) { + using Kernel = decltype(_k); + using scalar_t = typename Kernel::scalar_t; + (void)_k; + + if (kernel_launched) { + return; + } + + if (dim_value > Kernel::kMaxK || dim > Kernel::kMaxK) { + return; + } + + if (natten::flatten(query_tile_shape) != Kernel::kQueriesPerBlock || + natten::flatten(key_tile_shape) != Kernel::kKeysPerBlock) { + return; + } + + // Alignment + if ((dim % Kernel::kAlignmentQ) || (dim % Kernel::kAlignmentK) || + (dim_value % Kernel::kAlignmentV)) { + return; + } + // Uses too much smem + size_t smem_bytes = sizeof(typename Kernel::SharedStorage); + if (smem_bytes > max_smem) { + return; + } + + kernel_launched = true; + + typename Kernel::Params p; + p.query_ptr = (scalar_t*)query_ptr; + p.key_ptr = (scalar_t*)key_ptr; + p.value_ptr = (scalar_t*)value_ptr; + p.logsumexp_ptr = compute_logsumexp + ? (typename Kernel::lse_scalar_t*)logsumexp_ptr + : nullptr; + + // void* accum_ptr = nullptr; + if (Kernel::kNeedsOutputAccumulatorBuffer) { + using AccumType = typename Kernel::output_accum_t; + // TODO: let PT reserve space for accum_ptr and LSE. + // cudaMalloc-ing can break CUDA graphs, and probably does not + // guarantee the space is allocated on the same device (when multiple + // devices are present). + // cudaMallocAsync takes in the stream, which --should-- solve both + // issues. + // cudaMalloc(&accum_ptr, batch_size * natten::flatten(spatial_extent) * + // heads * dim_value * sizeof(AccumType)); + int64_t workspace_size_bytes = batch_size * + natten::flatten(spatial_extent) * heads * dim_value * + sizeof(AccumType); + void* accum_ptr = nullptr; + alloc_bytes(&accum_ptr, workspace_size_bytes, /* zero_fill = */ false); + // NATTEN_CUDA_CHECK(cudaMallocAsync( + // &accum_ptr, + // batch_size * natten::flatten(spatial_extent) * heads * dim_value * + // sizeof(AccumType), + // stream)); + p.output_accum_ptr = (AccumType*)accum_ptr; + } else { + p.output_accum_ptr = nullptr; + } + p.output_ptr = (typename Kernel::output_t*)out_ptr; + + p.num_heads = heads; + p.head_dim = dim; + p.head_dim_value = dim_value; + p.num_queries = tuple_to_na_dim(spatial_extent); + p.num_batches = batch_size; + + p.scale = attn_scale; + + p.kernel_size = tuple_to_na_dim(kernel_size); + p.stride = tuple_to_na_dim(stride); + p.dilation = tuple_to_na_dim(dilation); + + p.query_tile_shape = tuple_to_na_dim(query_tile_shape); + p.key_tile_shape = tuple_to_na_dim(key_tile_shape); + + if (smem_bytes > 0xc000) { + auto err = cudaFuncSetAttribute( + kernel_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes); + NATTEN_CHECK( + err != cudaErrorInvalidValue, + "This GPU does not have enough shared-memory."); + NATTEN_CUDA_CHECK(err); + } + auto blocks = p.getBlocksGrid(); + Kernel::check_supported(p); + kernel_fn<<>>(p); + }; + + DISPATCH_FNA_FORWARD_KERNEL(kRank, cc, dtype, is_causal, launchKernel); + NATTEN_CHECK( + kernel_launched, + "Could not find a compatible fused neighborhood attention kernel."); + NATTEN_CUDA_CHECK(cudaGetLastError()); +} + +} // namespace fna +} // namespace cuda +} // namespace natten diff --git a/natten/csrc/include/natten/cuda/fna/gemm/custom_mma.h b/natten/csrc/include/natten/cuda/fna/gemm/custom_mma.h new file mode 100644 index 00000000..a8b5a01c --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna/gemm/custom_mma.h @@ -0,0 +1,271 @@ +/* + * Copied from xFormers (https://github.com/facebookresearch/xformers/) and + * edited + * + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * BSD 3-Clause License + * + * 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 names of Facebook, Deepmind Technologies, NYU, NEC + * Laboratories America and IDIAP Research Institute 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 OWNER 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 +#include + +#include +#include + +namespace natten { +namespace cuda { +namespace fna { + +template +struct MakeCustomMma; + +template < + typename Shape, + typename IteratorA, + typename SmemIteratorA, + cutlass::arch::CacheOperation::Kind CacheOpA, + typename IteratorB, + typename SmemIteratorB, + cutlass::arch::CacheOperation::Kind CacheOpB, + typename ElementC, + typename LayoutC, + typename Policy, + int Stages, + cutlass::gemm::SharedMemoryClearOption SharedMemoryClear, + int kMaxK> +struct MakeCustomMma< + cutlass::gemm::threadblock::MmaMultistage< + Shape, + IteratorA, + SmemIteratorA, + CacheOpA, + IteratorB, + SmemIteratorB, + CacheOpB, + ElementC, + LayoutC, + Policy, + Stages, + SharedMemoryClear>, + kMaxK> { + // Reduce the number of stages if we don't need that many + static int constexpr kStages = + kMaxK == cutlass::platform::numeric_limits::max() + ? Stages + : cutlass::const_min( + Stages, + (kMaxK + int(Shape::kK) - 1) / int(Shape::kK)); + using Mma = cutlass::gemm::threadblock::CustomMmaMultistage< + Shape, + IteratorA, + SmemIteratorA, + CacheOpA, + IteratorB, + SmemIteratorB, + CacheOpB, + ElementC, + LayoutC, + Policy, + kStages, + SharedMemoryClear, + kMaxK>; +}; + +template < + typename Shape, + typename IteratorA, + typename SmemIteratorA, + typename IteratorB, + typename SmemIteratorB, + typename ElementC, + typename LayoutC, + typename Policy, + int kMaxK> +struct MakeCustomMma< + cutlass::gemm::threadblock::MmaPipelined< + Shape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + Policy>, + kMaxK> { + using Mma = cutlass::gemm::threadblock::CustomMmaPipelined< + Shape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + Policy>; +}; + +template +struct MakeCustomMmaAndReplaceIterators; + +template < + int NADim, + typename Shape, + typename IteratorA, + typename SmemIteratorA, + cutlass::arch::CacheOperation::Kind CacheOpA, + typename IteratorB, + typename SmemIteratorB, + cutlass::arch::CacheOperation::Kind CacheOpB, + typename ElementC, + typename LayoutC, + typename Policy, + int Stages, + cutlass::gemm::SharedMemoryClearOption SharedMemoryClear, + int kMaxK> +struct MakeCustomMmaAndReplaceIterators< + NADim, + cutlass::gemm::threadblock::MmaMultistage< + Shape, + IteratorA, + SmemIteratorA, + CacheOpA, + IteratorB, + SmemIteratorB, + CacheOpB, + ElementC, + LayoutC, + Policy, + Stages, + SharedMemoryClear>, + kMaxK> { + using NewIteratorA = + cutlass::transform::threadblock::CustomPredicatedTileAccessIterator< + NADim, + typename IteratorA::Shape, + typename IteratorA::Element, + typename IteratorA::Layout, + IteratorA::kAdvanceRank, + typename IteratorA::ThreadMap, + typename IteratorA::AccessType>; + + using NewIteratorB = + cutlass::transform::threadblock::CustomPredicatedTileAccessIterator< + NADim, + typename IteratorB::Shape, + typename IteratorB::Element, + typename IteratorB::Layout, + IteratorB::kAdvanceRank, + typename IteratorB::ThreadMap, + typename IteratorB::AccessType>; + // Reduce the number of stages if we don't need that many + static int constexpr kStages = + kMaxK == cutlass::platform::numeric_limits::max() + ? Stages + : cutlass::const_min( + Stages, + (kMaxK + int(Shape::kK) - 1) / int(Shape::kK)); + using Mma = cutlass::gemm::threadblock::CustomMmaMultistage< + Shape, + NewIteratorA, + SmemIteratorA, + CacheOpA, + NewIteratorB, + SmemIteratorB, + CacheOpB, + ElementC, + LayoutC, + Policy, + kStages, + SharedMemoryClear, + kMaxK>; +}; + +template < + int NADim, + typename Shape, + typename IteratorA, + typename SmemIteratorA, + typename IteratorB, + typename SmemIteratorB, + typename ElementC, + typename LayoutC, + typename Policy, + int kMaxK> +struct MakeCustomMmaAndReplaceIterators< + NADim, + cutlass::gemm::threadblock::MmaPipelined< + Shape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + Policy>, + kMaxK> { + using NewIteratorA = + cutlass::transform::threadblock::CustomPredicatedTileIterator< + NADim, + typename IteratorA::Shape, + typename IteratorA::Element, + typename IteratorA::Layout, + IteratorA::kAdvanceRank, + typename IteratorA::ThreadMap, + IteratorA::AccessType::kElements>; + + using NewIteratorB = + cutlass::transform::threadblock::CustomPredicatedTileIterator< + NADim, + typename IteratorB::Shape, + typename IteratorB::Element, + typename IteratorB::Layout, + IteratorB::kAdvanceRank, + typename IteratorB::ThreadMap, + IteratorB::AccessType::kElements>; + using Mma = cutlass::gemm::threadblock::CustomMmaPipelined< + Shape, + NewIteratorA, + SmemIteratorA, + NewIteratorB, + SmemIteratorB, + ElementC, + LayoutC, + Policy>; +}; + +} // namespace fna +} // namespace cuda +} // namespace natten diff --git a/natten/csrc/include/natten/cuda/fna/gemm/custom_mma_base.h b/natten/csrc/include/natten/cuda/fna/gemm/custom_mma_base.h new file mode 100644 index 00000000..3da8befc --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna/gemm/custom_mma_base.h @@ -0,0 +1,189 @@ +/* + * Copied from xFormers (https://github.com/facebookresearch/xformers/) + * + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + */ +/*************************************************************************************************** + * Copyright (c) 2017 - 2026 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 Template for a double-buffered threadblock-scoped GEMM kernel. +*/ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace gemm { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Used for partial specialization + typename Enable = bool> +class CustomMmaBase { + public: + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + ///< Policy describing tuning details + using Policy = Policy_; + + // + // Dependent types + // + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Shape describing the overall GEMM computed from shared memory + /// by each warp. + using WarpGemm = typename Policy::Operator::Shape; + + /// Shape describing the number of warps filling the CTA + using WarpCount = GemmShape< + Shape::kM / WarpGemm::kM, + Shape::kN / WarpGemm::kN, + Shape::kK / WarpGemm::kK>; + + /// Number of warp-level GEMM oeprations + static int const kWarpGemmIterations = + (WarpGemm::kK / Operator::Policy::MmaShape::kK); + + /// Number of stages + static int const kStages = Stages; + + // + // Nested structs + // + + /// Shared storage object needed by threadblock-scoped GEMM + template + struct OperandSharedStorage { + AlignedBuffer buffer; + using TensorRef = TensorRef; + + CUTLASS_DEVICE + static OperandLayout Layout() { + return OperandLayout::packed({OperandShape::kRow, OperandShape::kColumn}); + } + + /// Returns a TensorRef to the operand + CUTLASS_HOST_DEVICE + TensorRef ref() { + return TensorRef{buffer.data(), Layout()}; + } + }; + + /// Shape of the A matrix operand in shared memory + using ShapeA = MatrixShape< + Shape::kM + Policy::SmemPaddingA::kRow, + Shape::kK * kStages + Policy::SmemPaddingA::kColumn>; + + /// Shape of the B matrix operand in shared memory + using ShapeB = MatrixShape< + Shape::kK * kStages + Policy::SmemPaddingB::kRow, + Shape::kN + Policy::SmemPaddingB::kColumn>; + + using SharedStorageA = OperandSharedStorage< + typename Operator::ElementA, + ShapeA, + typename Operator::LayoutA>; + using SharedStorageB = OperandSharedStorage< + typename Operator::ElementB, + ShapeB, + typename Operator::LayoutB>; + using TensorRefA = typename SharedStorageA::TensorRef; + using TensorRefB = typename SharedStorageB::TensorRef; + + struct SharedStorage { + /// Buffer for A operand + SharedStorageA operand_A; + + /// Buffer for B operand + SharedStorageB operand_B; + }; + + protected: + // + // Data members + // + + /// Iterator to load a warp-scoped tile of A operand from shared memory + typename Operator::IteratorA warp_tile_iterator_A_; + + /// Iterator to load a warp-scoped tile of B operand from shared memory + typename Operator::IteratorB warp_tile_iterator_B_; + + public: + /// Construct from tensor references + CUTLASS_DEVICE + CustomMmaBase( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + SharedStorageA& shared_storageA, + SharedStorageB& shared_storageB, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx) + : warp_tile_iterator_A_(shared_storageA.ref(), lane_idx), + warp_tile_iterator_B_(shared_storageB.ref(), lane_idx) {} +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fna/gemm/custom_mma_multistage.h b/natten/csrc/include/natten/cuda/fna/gemm/custom_mma_multistage.h new file mode 100644 index 00000000..509c2053 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna/gemm/custom_mma_multistage.h @@ -0,0 +1,765 @@ +/* + * Copied from xFormers (https://github.com/facebookresearch/xformers/) + * + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + */ +/*************************************************************************************************** + * Copyright (c) 2017 - 2026 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 Template for a double-buffered threadblock-scoped GEMM kernel. +*/ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace gemm { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + /// Cache operation for operand A + cutlass::arch::CacheOperation::Kind CacheOpA, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Cache operation for operand B + cutlass::arch::CacheOperation::Kind CacheOpB, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone, + /// Upper boundon the K dimension + int kMaxK = cutlass::platform::numeric_limits::max(), + /// Used for partial specialization + typename Enable = bool> +class CustomMmaMultistage : public CustomMmaBase { + public: + ///< Base class + using Base = CustomMmaBase; + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + ///< Iterates over tiles of A operand in global memory + using IteratorA = IteratorA_; + ///< Iterates over tiles of B operand in global memory + using IteratorB = IteratorB_; + ///< Data type of accumulator matrix + using ElementC = ElementC_; + ///< Layout of accumulator matrix + using LayoutC = LayoutC_; + ///< Policy describing tuning details + using Policy = Policy_; + + using SmemIteratorA = SmemIteratorA_; + using SmemIteratorB = SmemIteratorB_; + + static cutlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA; + static cutlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB; + + // + // Dependent types + // + + /// Fragment of accumulator tile + using FragmentC = typename Policy::Operator::FragmentC; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Minimum architecture is Sm80 to support cp.async + using ArchTag = arch::Sm80; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = Operator::kTransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = Operator::kTransformB; + + /// Internal structure exposed for introspection. + struct Detail { + static_assert( + Base::kWarpGemmIterations > 1, + "The pipelined structure requires at least two warp-level " + "GEMM operations."); + + /// Number of cp.async instructions to load one stage of operand A + static int const AsyncCopyIterationsPerStageA = + IteratorA::ThreadMap::Iterations::kCount; + + /// Number of cp.async instructions to load one stage of operand B + static int const AsyncCopyIterationsPerStageB = + IteratorB::ThreadMap::Iterations::kCount; + + /// Number of stages + static int const kStages = Stages; + + /// Number of cp.async instructions to load on group of operand A + static int const kAccessesPerGroupA = + (AsyncCopyIterationsPerStageA + Base::kWarpGemmIterations - 1) / + Base::kWarpGemmIterations; + + /// Number of cp.async instructions to load on group of operand B + static int const kAccessesPerGroupB = + (AsyncCopyIterationsPerStageB + Base::kWarpGemmIterations - 1) / + Base::kWarpGemmIterations; + }; + + static bool const kSmemContainsEntireMat = kMaxK <= Shape::kK * Stages; + static constexpr int kNumStagesConcurrentLoad = + kSmemContainsEntireMat ? Stages : Stages - 1; + + private: + using WarpLoadedFragmentA = typename Operator::FragmentA; + using WarpLoadedFragmentB = typename Operator::FragmentB; + using WarpTransformedFragmentA = typename Operator::TransformedFragmentA; + using WarpTransformedFragmentB = typename Operator::TransformedFragmentB; + + private: + // + // Data members + // + + /// Iterator to write threadblock-scoped tile of A operand to shared memory + SmemIteratorA smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB smem_iterator_B_; + + bool prologue_done_; + + // Set to `True` to ensure the accumulator will be zero outside the GEMM + // footprint + bool zero_outside_bounds_; + + public: + /// Construct from tensor references + CUTLASS_DEVICE + CustomMmaMultistage( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + typename Base::SharedStorageA& shared_storageA, + typename Base::SharedStorageB& shared_storageB, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx) + : Base(shared_storageA, shared_storageB, thread_idx, warp_idx, lane_idx), + smem_iterator_A_(shared_storageA.ref(), thread_idx), + smem_iterator_B_(shared_storageB.ref(), thread_idx), + prologue_done_(false), + zero_outside_bounds_(false) { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + + int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; + int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset( + {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); + } + CUTLASS_DEVICE + CustomMmaMultistage( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + typename Base::SharedStorage& st, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx) + : CustomMmaMultistage( + st.operand_A, + st.operand_B, + thread_idx, + warp_idx, + lane_idx) {} + + CUTLASS_DEVICE + void set_prologue_done(bool value) { + prologue_done_ = value; + } + + CUTLASS_DEVICE + void set_zero_outside_bounds(bool value) { + zero_outside_bounds_ = value; + } + + template + CUTLASS_DEVICE static void prologue( + typename Base::SharedStorage& shared_storage, + ///< iterator over A operand in global memory + IteratorA iterator_A, + ///< iterator over B operand in global memory + IteratorB iterator_B, + int thread_idx, + int problem_size_k) { + prologue( + shared_storage.operand_A, + shared_storage.operand_B, + iterator_A, + iterator_B, + thread_idx, + problem_size_k); + } + + template + CUTLASS_DEVICE static void prologue( + typename Base::SharedStorageA& shared_storageA, + typename Base::SharedStorageB& shared_storageB, + ///< iterator over A operand in global memory + IteratorA iterator_A, + ///< iterator over B operand in global memory + IteratorB iterator_B, + int thread_idx, + int problem_size_k) { + SmemIteratorA smem_iterator_A(shared_storageA.ref(), thread_idx); + SmemIteratorB smem_iterator_B(shared_storageB.ref(), thread_idx); + int32_t iter = (problem_size_k + Base::Shape::kK - 1) / Base::Shape::kK; + _prologue( + iterator_A, iterator_B, iter, smem_iterator_A, smem_iterator_B); + } + + CUTLASS_DEVICE + void copy_tiles_and_advance( + IteratorA& iterator_A, + IteratorB& iterator_B, + int group_start_A = 0, + int group_start_B = 0) { + iterator_A.set_iteration_index( + group_start_A * IteratorA::kAccessesPerVector); + this->smem_iterator_A_.set_iteration_index(group_start_A); + + // Async Copy for operand A + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupA; ++j) { + if (group_start_A + j < Detail::AsyncCopyIterationsPerStageA) { + typename IteratorA::AccessType* dst_ptr = + reinterpret_cast( + this->smem_iterator_A_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / + IteratorA::kAccessesPerVector / 8; + + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { + auto gmem_ptr = iterator_A.get(); + + if (zero_outside_bounds_ || + SharedMemoryClear == SharedMemoryClearOption::kZfill) { + cutlass::arch::cp_async_zfill( + dst_ptr + v, gmem_ptr, iterator_A.valid()); + } else { + cutlass::arch::cp_async( + dst_ptr + v, gmem_ptr, iterator_A.valid()); + } + + ++iterator_A; + } + + ++this->smem_iterator_A_; + } + } + + iterator_B.set_iteration_index( + group_start_B * IteratorB::kAccessesPerVector); + this->smem_iterator_B_.set_iteration_index(group_start_B); + + // Async Copy for operand B + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupB; ++j) { + if (group_start_B + j < Detail::AsyncCopyIterationsPerStageB) { + typename IteratorB::AccessType* dst_ptr = + reinterpret_cast( + this->smem_iterator_B_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / + IteratorB::kAccessesPerVector / 8; + + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) { + auto gmem_ptr = iterator_B.get(); + + if (zero_outside_bounds_ || + SharedMemoryClear == SharedMemoryClearOption::kZfill) { + cutlass::arch::cp_async_zfill( + dst_ptr + v, gmem_ptr, iterator_B.valid()); + } else { + cutlass::arch::cp_async( + dst_ptr + v, gmem_ptr, iterator_B.valid()); + } + + ++iterator_B; + } + ++this->smem_iterator_B_; + } + } + } + + template + CUTLASS_DEVICE static void _prologue( + IteratorA& iterator_A, + IteratorB& iterator_B, + int32_t& gemm_k_iterations, + SmemIteratorA& smem_iterator_A_, + SmemIteratorB& smem_iterator_B_) { + // Issue several complete stages + CUTLASS_PRAGMA_UNROLL + for (int stage = 0; stage < kNumStagesConcurrentLoad; + ++stage, --gemm_k_iterations) { + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + + iterator_A.set_iteration_index(0); + smem_iterator_A_.set_iteration_index(0); + + // Async Copy for operand A + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) { + typename IteratorA::AccessType* dst_ptr = + reinterpret_cast( + smem_iterator_A_.get()); + + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { + int const kSrcBytes = + sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / + IteratorA::kAccessesPerVector / 8; + + int src_bytes = (iterator_A.valid() ? kSrcBytes : 0); + + if (kLoadA) { + cutlass::arch::cp_async_zfill( + dst_ptr + v, iterator_A.get(), iterator_A.valid()); + } + + ++iterator_A; + } + + ++smem_iterator_A_; + } + + iterator_B.set_iteration_index(0); + smem_iterator_B_.set_iteration_index(0); + + // Async Copy for operand B + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) { + typename IteratorB::AccessType* dst_ptr = + reinterpret_cast( + smem_iterator_B_.get()); + + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) { + int const kSrcBytes = + sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / + IteratorB::kAccessesPerVector / 8; + + if (kLoadB) { + cutlass::arch::cp_async_zfill( + dst_ptr + v, iterator_B.get(), iterator_B.valid()); + } + + ++iterator_B; + } + + ++smem_iterator_B_; + } + + // Move to the next stage + iterator_A.add_tile_offset({0, 1}); + iterator_B.add_tile_offset({1, 0}); + + smem_iterator_A_.add_tile_offset({0, 1}); + smem_iterator_B_.add_tile_offset({1, 0}); + + // Defines the boundary of a stage of cp.async. + cutlass::arch::cp_async_fence(); + } + } + + /// Perform a threadblock-scoped matrix multiply-accumulate + CUTLASS_DEVICE + void operator()( + ///< problem size of GEMM + int gemm_k_iterations, + ///< destination accumulator tile + FragmentC& accum, + ///< iterator over A operand in global memory + IteratorA iterator_A, + ///< iterator over B operand in global memory + IteratorB iterator_B, + ///< initial value of accumulator + FragmentC const& src_accum) { + // + // Prologue + // + + if (!prologue_done_) { + _prologue( + iterator_A, + iterator_B, + gemm_k_iterations, + smem_iterator_A_, + smem_iterator_B_); + } else if (!kSmemContainsEntireMat) { + _prologue( + iterator_A, + iterator_B, + gemm_k_iterations, + smem_iterator_A_, + smem_iterator_B_); + } else { + gemm_k_iterations -= kNumStagesConcurrentLoad; + } + + // Perform accumulation in the 'd' output operand + accum = src_accum; + + // + // Clear the remaining tiles of SMEM. This is a functional requirement for + // some kernels so that all accumulator elements outside the GEMM footprint + // are zero. + // + + if (SharedMemoryClear == SharedMemoryClearOption::kClearLastStage) { + /// Iterator to write threadblock-scoped tile of A operand to shared + /// memory + SmemIteratorA last_smem_iterator_A(this->smem_iterator_A_); + + typename IteratorA::AccessType zero_A; + zero_A.clear(); + + last_smem_iterator_A.set_iteration_index(0); + + // Async Copy for operand A + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) { + typename IteratorA::AccessType* dst_ptr = + reinterpret_cast( + last_smem_iterator_A.get()); + + *dst_ptr = zero_A; + + ++last_smem_iterator_A; + } + + /// Iterator to write threadblock-scoped tile of B operand to shared + /// memory + SmemIteratorB last_smem_iterator_B(this->smem_iterator_B_); + typename IteratorB::AccessType zero_B; + + zero_B.clear(); + last_smem_iterator_B.set_iteration_index(0); + + // Async Copy for operand B + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) { + typename IteratorB::AccessType* dst_ptr = + reinterpret_cast( + last_smem_iterator_B.get()); + + *dst_ptr = zero_B; + + ++last_smem_iterator_B; + } + } + + // Waits until kStages-2 stages have committed. + cutlass::arch::cp_async_wait(); + __syncthreads(); + + // Pair of fragments used to overlap shared memory loads and math + // instructions + WarpLoadedFragmentA warp_loaded_frag_A[2]; + WarpLoadedFragmentB warp_loaded_frag_B[2]; + WarpTransformedFragmentA warp_transformed_frag_A[2]; + WarpTransformedFragmentB warp_transformed_frag_B[2]; + + Operator warp_mma; + + this->warp_tile_iterator_A_.set_kgroup_index(0); + this->warp_tile_iterator_B_.set_kgroup_index(0); + + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[0]); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B[0]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + + int smem_write_stage_idx = Base::kStages - 1; + int smem_read_stage_idx = 0; + + warp_mma.transform( + warp_transformed_frag_A[0], + warp_transformed_frag_B[0], + warp_loaded_frag_A[0], + warp_loaded_frag_B[0]); + + // tf32x3 kernels use staging accumulation. warp_mma uses a temporary + // accumulator and this temporary accumulator is added to the final + // accumulator once in every mainloop iteration. + plus plus_accum; + + FragmentC tmp_accum; + + if (platform::is_same< + typename Operator::MathOperator, + arch::OpMultiplyAddFastF32>::value || + platform::is_same< + typename Operator::MathOperator, + arch::OpMultiplyAddComplexFastF32>::value) { + tmp_accum.clear(); + } + + // + // Mainloop + // + + CUTLASS_GEMM_LOOP + for (; gemm_k_iterations > (-kNumStagesConcurrentLoad);) { + // + // Loop over GEMM K dimension + // + + // Computes a warp-level GEMM on data held in shared memory + // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate + CUTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; + ++warp_mma_k) { + // Load warp-level tiles from shared memory, wrapping to k offset if + // this is the last group as the case may be. + + this->warp_tile_iterator_A_.set_kgroup_index( + (warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_B_.set_kgroup_index( + (warp_mma_k + 1) % Base::kWarpGemmIterations); + + // In case of a non-circular buffer ("kSmemContainsEntireMat") + // make sure we don't load out of bounds data. + if (!kSmemContainsEntireMat || + gemm_k_iterations > (-kNumStagesConcurrentLoad) || + warp_mma_k < Base::kWarpGemmIterations - 1) { + this->warp_tile_iterator_A_.load( + warp_loaded_frag_A[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_B_.load( + warp_loaded_frag_B[(warp_mma_k + 1) % 2]); + } + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + if (warp_mma_k > 0) + warp_mma.transform( + warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + warp_loaded_frag_A[warp_mma_k % 2], + warp_loaded_frag_B[warp_mma_k % 2]); + + if (platform::is_same< + typename Operator::MathOperator, + arch::OpMultiplyAddFastF32>::value || + platform::is_same< + typename Operator::MathOperator, + arch::OpMultiplyAddComplexFastF32>::value) { + warp_mma( + tmp_accum, + warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + tmp_accum); + + if (warp_mma_k == 0) { + accum = plus_accum(accum, tmp_accum); + tmp_accum.clear(); + } + } else { + warp_mma( + accum, + warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + accum); + } + + // Issue global->shared copies for the this stage + if (!kSmemContainsEntireMat && + warp_mma_k < Base::kWarpGemmIterations - 1) { + int group_start_iteration_A, group_start_iteration_B; + + group_start_iteration_A = warp_mma_k * Detail::kAccessesPerGroupA; + group_start_iteration_B = warp_mma_k * Detail::kAccessesPerGroupB; + + copy_tiles_and_advance( + iterator_A, + iterator_B, + group_start_iteration_A, + group_start_iteration_B); + } + + if (warp_mma_k + 2 == Base::kWarpGemmIterations) { + if (!kSmemContainsEntireMat) { + int group_start_iteration_A, group_start_iteration_B; + group_start_iteration_A = + (warp_mma_k + 1) * Detail::kAccessesPerGroupA; + group_start_iteration_B = + (warp_mma_k + 1) * Detail::kAccessesPerGroupB; + + copy_tiles_and_advance( + iterator_A, + iterator_B, + group_start_iteration_A, + group_start_iteration_B); + } + + // Inserts a memory fence between stages of cp.async instructions. + cutlass::arch::cp_async_fence(); + + // Waits until kStages-2 stages have committed. + cutlass::arch::cp_async_wait(); + __syncthreads(); + + // Move to the next stage + iterator_A.add_tile_offset({0, 1}); + iterator_B.add_tile_offset({1, 0}); + + this->smem_iterator_A_.add_tile_offset({0, 1}); + this->smem_iterator_B_.add_tile_offset({1, 0}); + + // Add negative offsets to return iterators to the 'start' of the + // circular buffer in shared memory + if (smem_write_stage_idx == (Base::kStages - 1)) { + this->smem_iterator_A_.add_tile_offset({0, -Base::kStages}); + this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); + smem_write_stage_idx = 0; + } else { + ++smem_write_stage_idx; + } + + if (!kSmemContainsEntireMat && + smem_read_stage_idx == (Base::kStages - 1)) { + this->warp_tile_iterator_A_.add_tile_offset( + {0, + -Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations}); + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations, + 0}); + smem_read_stage_idx = 0; + } else { + ++smem_read_stage_idx; + } + + --gemm_k_iterations; + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + } + + // Do any conversions feeding the first stage at the end of the loop so + // we can start right away on mma instructions + if (warp_mma_k + 1 == Base::kWarpGemmIterations) + warp_mma.transform( + warp_transformed_frag_A[(warp_mma_k + 1) % 2], + warp_transformed_frag_B[(warp_mma_k + 1) % 2], + warp_loaded_frag_A[(warp_mma_k + 1) % 2], + warp_loaded_frag_B[(warp_mma_k + 1) % 2]); + } + } + + if (platform::is_same< + typename Operator::MathOperator, + arch::OpMultiplyAddFastF32>::value || + platform::is_same< + typename Operator::MathOperator, + arch::OpMultiplyAddComplexFastF32>::value) { + accum = plus_accum(accum, tmp_accum); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fna/gemm/custom_mma_pipelined.h b/natten/csrc/include/natten/cuda/fna/gemm/custom_mma_pipelined.h new file mode 100644 index 00000000..c4c892e0 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna/gemm/custom_mma_pipelined.h @@ -0,0 +1,408 @@ +/* + * Copied from xFormers (https://github.com/facebookresearch/xformers/) + * + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + */ +/*************************************************************************************************** + * Copyright (c) 2017 - 2026 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 Template for a double-buffered threadblock-scoped GEMM kernel. +*/ + +#pragma once + +#include +#include +#include +#include + +#include +#include + +#include +#include + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace gemm { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Transformation applied to A operand + typename TransformA_ = NumericArrayConverter< + typename SmemIteratorA_::Element, + typename IteratorA_::Element, + IteratorA_::Fragment::kElements>, + /// + /// Transformation applied to B operand + typename TransformB_ = NumericArrayConverter< + typename SmemIteratorB_::Element, + typename IteratorB_::Element, + IteratorB_::Fragment::kElements>, + /// Used for partial specialization + typename Enable = bool> +class CustomMmaPipelined : public CustomMmaBase { + public: + ///< Base class + using Base = CustomMmaBase; + + using Shape = + Shape_; ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using IteratorA = + IteratorA_; ///< Iterates over tiles of A operand in global memory + using IteratorB = + IteratorB_; ///< Iterates over tiles of B operand in global memory + using ElementC = ElementC_; ///< Data type of accumulator matrix + using LayoutC = LayoutC_; ///< Layout of accumulator matrix + using Policy = Policy_; ///< Policy describing tuning details + + using SmemIteratorA = SmemIteratorA_; + using SmemIteratorB = SmemIteratorB_; + + using TransformA = TransformA_; + using TransformB = TransformB_; + + // + // Dependent types + // + + /// Fragment of operand A loaded from global memory + using FragmentA = typename IteratorA::Fragment; + + /// Fragment of operand B loaded from global memory + using FragmentB = typename IteratorB::Fragment; + + /// Fragment of accumulator tile + using FragmentC = typename Policy::Operator::FragmentC; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Obtain the arch tag from the warp-level operator + using ArchTag = typename Policy::Operator::ArchTag; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = Operator::kTransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = Operator::kTransformB; + + // staticaly assert kStages for MmaPipelined is two (Double-buffered pipeline) + static_assert( + (Base::kStages == 2), + "MmaPipelined requires kStages set to value 2"); + + static bool const kSmemContainsEntireMat = false; + + private: + using WarpFragmentA = typename Operator::FragmentA; + using WarpFragmentB = typename Operator::FragmentB; + + protected: + /// Iterator to write threadblock-scoped tile of A operand to shared memory + SmemIteratorA smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB smem_iterator_B_; + + public: + /// Construct from tensor references + CUTLASS_DEVICE + CustomMmaPipelined( + typename Base::SharedStorageA& shared_storageA, + typename Base::SharedStorageB& shared_storageB, + int thread_idx, ///< ID within the threadblock + int warp_idx, ///< ID of warp + int lane_idx ///< ID of each thread within a warp + ) + : Base(shared_storageA, shared_storageB, thread_idx, warp_idx, lane_idx), + smem_iterator_A_(shared_storageA.ref(), thread_idx), + smem_iterator_B_(shared_storageB.ref(), thread_idx) { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + + int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; + int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset( + {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); + } + CUTLASS_DEVICE + CustomMmaPipelined( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + typename Base::SharedStorage& st, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx) + : CustomMmaPipelined( + st.operand_A, + st.operand_B, + thread_idx, + warp_idx, + lane_idx) {} + + CUTLASS_DEVICE + void set_prologue_done(bool value) { + // NOT IMPLEMENTED FOR PIPELINED + } + + CUTLASS_DEVICE + void set_zero_outside_bounds(bool value) { + // NOT NEEDED FOR PIPELINED + // shared memory will always be zero-filled + } + + template + CUTLASS_DEVICE static void prologue( + typename Base::SharedStorage& shared_storage, + ///< iterator over A operand in global memory + IteratorA iterator_A, + ///< iterator over B operand in global memory + IteratorB iterator_B, + int thread_idx, + int problem_size_k) { + prologue( + shared_storage.operand_A, + shared_storage.operand_B, + iterator_A, + iterator_B, + thread_idx, + problem_size_k); + } + + template + CUTLASS_DEVICE static void prologue( + typename Base::SharedStorageA& shared_storageA, + typename Base::SharedStorageB& shared_storageB, + ///< iterator over A operand in global memory + IteratorA iterator_A, + ///< iterator over B operand in global memory + IteratorB iterator_B, + int thread_idx, + int problem_size_k) { + // NOT IMPLEMENTED FOR PIPELINED + } + + /// Perform a threadblock-scoped matrix multiply-accumulate + CUTLASS_DEVICE + void operator()( + int gemm_k_iterations, ///< number of iterations of the mainloop + FragmentC& accum, ///< destination accumulator tile + IteratorA iterator_A, ///< iterator over A operand in global memory + IteratorB iterator_B, ///< iterator over B operand in global memory + FragmentC const& src_accum, ///< source accumulator tile + TransformA transform_A = + TransformA(), ///< transformation applied to A fragment + TransformB transform_B = + TransformB()) { ///< transformation applied to B fragment + + // + // Prologue + // + + // Perform accumulation in the 'd' output operand + accum = src_accum; + + FragmentA tb_frag_A; + FragmentB tb_frag_B; + + tb_frag_A.clear(); + tb_frag_B.clear(); + + // The last kblock is loaded in the prolog + iterator_A.load(tb_frag_A); + iterator_B.load(tb_frag_B); + + ++iterator_A; + ++iterator_B; + + this->smem_iterator_A_.store(transform_A(tb_frag_A)); + this->smem_iterator_B_.store(transform_B(tb_frag_B)); + + ++this->smem_iterator_A_; + ++this->smem_iterator_B_; + + __syncthreads(); + + // Pair of fragments used to overlap shared memory loads and math + // instructions + WarpFragmentA warp_frag_A[2]; + WarpFragmentB warp_frag_B[2]; + + this->warp_tile_iterator_A_.set_kgroup_index(0); + this->warp_tile_iterator_B_.set_kgroup_index(0); + + this->warp_tile_iterator_A_.load(warp_frag_A[0]); + this->warp_tile_iterator_B_.load(warp_frag_B[0]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + Operator warp_mma; + + int smem_write_stage_idx = 1; + + // Avoid reading out of bounds + iterator_A.clear_mask(gemm_k_iterations <= 1); + iterator_B.clear_mask(gemm_k_iterations <= 1); + + // Issue loads during the first warp-level matrix multiply-add *AFTER* + // issuing shared memory loads (which have the tightest latency + // requirement). + + // + // Mainloop + // + + // Note: The main loop does not support Base::kWarpGemmIterations == 2. + CUTLASS_GEMM_LOOP + for (; gemm_k_iterations > 0; --gemm_k_iterations) { + // + // Loop over GEMM K dimension + // + + CUTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; + ++warp_mma_k) { + // Load warp-level tiles from shared memory, wrapping to k offset if + // this is the last group as the case may be. + + if (warp_mma_k == Base::kWarpGemmIterations - 1) { + // Write fragments to shared memory + this->smem_iterator_A_.store(transform_A(tb_frag_A)); + + this->smem_iterator_B_.store(transform_B(tb_frag_B)); + + __syncthreads(); + + ++this->smem_iterator_A_; + ++this->smem_iterator_B_; + + // Add negative offsets to return iterators to the 'start' of the + // circular buffer in shared memory + if (smem_write_stage_idx == 1) { + this->smem_iterator_A_.add_tile_offset({0, -Base::kStages}); + this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); + } else { + this->warp_tile_iterator_A_.add_tile_offset( + {0, + -Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations}); + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations, + 0}); + } + + smem_write_stage_idx ^= 1; + } + + this->warp_tile_iterator_A_.set_kgroup_index( + (warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_B_.set_kgroup_index( + (warp_mma_k + 1) % Base::kWarpGemmIterations); + + this->warp_tile_iterator_A_.load(warp_frag_A[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_B_.load(warp_frag_B[(warp_mma_k + 1) % 2]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + if (warp_mma_k == 0) { + iterator_A.load(tb_frag_A); + iterator_B.load(tb_frag_B); + + ++iterator_A; + ++iterator_B; + + // Avoid reading out of bounds if this was the last loop iteration + iterator_A.clear_mask(gemm_k_iterations <= 2); + iterator_B.clear_mask(gemm_k_iterations <= 2); + } + + warp_mma( + accum, + warp_frag_A[warp_mma_k % 2], + warp_frag_B[warp_mma_k % 2], + accum); + } + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fna/gemm/find_default_mma.h b/natten/csrc/include/natten/cuda/fna/gemm/find_default_mma.h new file mode 100644 index 00000000..47f9ef21 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna/gemm/find_default_mma.h @@ -0,0 +1,272 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Fused Neighborhood Attention kernels are heavily based on the + * memory-efficient attention kernels from the xFormers project by Meta + * Platforms, Inc. + * + * Copyright (c) Facebook, Inc. and its affiliates + * + * BSD 3-Clause License + * + * 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 names of Facebook, Deepmind Technologies, NYU, NEC + * Laboratories America and IDIAP Research Institute 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 OWNER 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 Cutlass provides helper template functions to figure out the right + datastructures to instanciate to run a GEMM with various parameters (see + `cutlass/gemm/threadblock/default_mma.h`). However, due to template + instantiation priority rules, it will only create an MmaMultiStage with + kStages=3 (otherwise creates an MmePipelined - which is not compatible with + FastF32). kStages=3 uses too much shared memory and we want to use kStages=2, + so we just copy-pasted some code from `default_mma.h` and + `default_mma_core.h` files and wrapped this template to allow our usecase. + + This is really only for the FastF32 case - aka using TensorCores with fp32. +*/ + +#pragma once + +#include +#include +#include +#include +#include + +#include +#include + +#include + +namespace cutlass { +namespace gemm { +namespace threadblock { + +template < + /// Neighborhood attention spatial rank; must be 0 (self attention), 1, 2, + /// or 3 + int NADim, + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Layout type for C and D matrix operand + typename LayoutC, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation perfomed by GEMM + typename Operator, + typename Enable_ = void> +struct FindDefaultMma { + static_assert(NADim >= 1 && NADim < 4); + static constexpr bool AccumulatorsInRowMajor = false; + static constexpr SharedMemoryClearOption SharedMemoryClear = + SharedMemoryClearOption::kNone; + using DefaultMma_ = cutlass::gemm::threadblock::DefaultMma< + ElementA, + LayoutA, + kAlignmentA, + ElementB, + LayoutB, + kAlignmentB, + ElementAccumulator, + LayoutC, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + Stages, + Operator, + AccumulatorsInRowMajor, + SharedMemoryClear>; + struct DefaultMma : DefaultMma_ { + using MmaCore_ = typename DefaultMma_::MmaCore; + using IteratorA_ = typename DefaultMma_::IteratorA; + using IteratorB_ = typename DefaultMma_::IteratorB; + using ThreadblockMma_ = typename DefaultMma_::ThreadblockMma; + + // struct ThreadblockMma : ThreadblockMma_ { + using IteratorA = + cutlass::transform::threadblock::CustomPredicatedTileIterator< + NADim, + typename IteratorA_::Shape, + typename IteratorA_::Element, + typename IteratorA_::Layout, + IteratorA_::kAdvanceRank, + typename IteratorA_::ThreadMap, + kAlignmentA>; + + using IteratorB = + cutlass::transform::threadblock::CustomPredicatedTileIterator< + NADim, + typename IteratorB_::Shape, + typename IteratorB_::Element, + typename IteratorB_::Layout, + IteratorB_::kAdvanceRank, + typename IteratorB_::ThreadMap, + kAlignmentB>; + //}; + + // using IteratorA = typename ThreadblockMma::IteratorA; + // using IteratorB = typename ThreadblockMma::IteratorA; + + using ThreadblockMma = typename natten::cuda::fna:: + ReplaceMmaIterators::Mma; + }; +}; + +/// Specialization for sm80 / FastF32 / multistage with kStages=2 +template < + /// Neighborhood attention spatial rank; must be 0 (self attention), 1, 2, + /// or 3 + int NADim, + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + int kStages, + typename Operator> +struct FindDefaultMma< + NADim, + ElementA_, + LayoutA_, + kAlignmentA, + ElementB_, + LayoutB_, + kAlignmentB, + ElementAccumulator, + layout::RowMajor, + arch::OpClassTensorOp, + arch::Sm80, + ThreadblockShape, + WarpShape, + InstructionShape, + kStages, + Operator, + typename cutlass::platform::enable_if<(kAlignmentA > 1)>::type> { + static_assert(NADim >= 1 && NADim < 4); + using LayoutC = layout::RowMajor; + using OperatorClass = arch::OpClassTensorOp; + using ArchTag = arch::Sm80; + + using DefaultMma_ = cutlass::gemm::threadblock::DefaultMma< + ElementA_, + LayoutA_, + kAlignmentA, + ElementB_, + LayoutB_, + kAlignmentB, + ElementAccumulator, + LayoutC, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + 3, + Operator>; + struct DefaultMma : DefaultMma_ { + using MmaCore_ = typename DefaultMma_::MmaCore; + using IteratorA_ = typename DefaultMma_::IteratorA; + using IteratorB_ = typename DefaultMma_::IteratorB; + + using IteratorA = + cutlass::transform::threadblock::CustomPredicatedTileAccessIterator< + NADim, + typename IteratorA_::Shape, + typename IteratorA_::Element, + typename IteratorA_::Layout, + IteratorA_::kAdvanceRank, + typename IteratorA_::ThreadMap, + typename DefaultMma_::AccessTypeA>; + + using IteratorB = + cutlass::transform::threadblock::CustomPredicatedTileAccessIterator< + NADim, + typename IteratorB_::Shape, + typename IteratorB_::Element, + typename IteratorB_::Layout, + IteratorB_::kAdvanceRank, + typename IteratorB_::ThreadMap, + typename DefaultMma_::AccessTypeB>; + + // Define the threadblock-scoped multistage matrix multiply + using ThreadblockMma = cutlass::gemm::threadblock::MmaMultistage< + typename MmaCore_::Shape, + IteratorA, + typename MmaCore_::SmemIteratorA, + MmaCore_::kCacheOpA, + IteratorB, + typename MmaCore_::SmemIteratorB, + MmaCore_::kCacheOpB, + ElementAccumulator, + LayoutC, + typename MmaCore_::MmaPolicy, + kStages>; + }; +}; + +} // namespace threadblock +} // namespace gemm +} // namespace cutlass diff --git a/natten/csrc/include/natten/cuda/fna/gemm/mma_accum_lambda_iterator.h b/natten/csrc/include/natten/cuda/fna/gemm/mma_accum_lambda_iterator.h new file mode 100644 index 00000000..2fcf3f2a --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna/gemm/mma_accum_lambda_iterator.h @@ -0,0 +1,393 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Fused Neighborhood Attention kernels are heavily based on the + * memory-efficient attention kernels from the xFormers project by Meta + * Platforms, Inc. + * + * Copyright (c) Facebook, Inc. and its affiliates + * + * BSD 3-Clause License + * + * 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 names of Facebook, Deepmind Technologies, NYU, NEC + * Laboratories America and IDIAP Research Institute 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 OWNER 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 +#include +#include + +/* +TensorCores have different accumulator layouts. +This file provides a class to easily map the accumulator +i-th element with the corresponding matrix row/col. +*/ + +template +struct AccumLambdaIteratorSm80 { + static_assert( + cutlass::platform:: + is_same::value, + "only RowMajor is supported"); + + using Policy = typename T::Policy; + using InstructionShape = typename T::InstructionShape; + using OpDelta = typename T::OpDelta; + using Shape = typename T::Shape; + static int const kElementsPerAccess = InstructionShape::kN / 4; + static int const kRowsPerTile = 8; + static int const kAccumulatorRows = InstructionShape::kM / kRowsPerTile; + + static int const kRowIters = Policy::MmaIterations::kRow * kAccumulatorRows; + + static cutlass::MatrixCoord CUTLASS_DEVICE get_lane_offset( + int8_t lane_id, + int8_t warp_id, + typename T::TensorCoord const& tile_offset) { + int quad = (lane_id >> 2); + int lane_in_quad = (lane_id & 3); + return cutlass::MatrixCoord( + quad + tile_offset.row() * Shape::kRow, + lane_in_quad * kElementsPerAccess + + tile_offset.column() * Shape::kColumn); + } + + template + CUTLASS_DEVICE static void iterateRows( + cutlass::MatrixCoord& lane_offset, + FA beginRow, + FB op, + FC endRow) { + // See cutlass/gemm/warp/mma_tensor_op_tile_iterator.h + CUTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + CUTLASS_PRAGMA_UNROLL + for (int row = 0; row < kAccumulatorRows; ++row) { + int accum_m = mma_m * InstructionShape::kM * OpDelta::kRow + + row * kRowsPerTile + lane_offset.row(); + beginRow(accum_m); + + CUTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + int mma_accum_start = kAccumulatorRows * kElementsPerAccess * + (mma_n * Policy::MmaIterations::kRow + mma_m); + CUTLASS_PRAGMA_UNROLL + for (int col = 0; col < kElementsPerAccess; ++col) { + int accum_n = mma_n * InstructionShape::kN * OpDelta::kColumn + + col + lane_offset.column(); + int idx = mma_accum_start + row * kElementsPerAccess + col; + op(accum_m, accum_n, idx); + } + } + + endRow(accum_m); + } + } + } + + template + CUTLASS_DEVICE static bool reduceSameRow(int lane_id, DT& myValue, F fn) { + // In each warp, 4 threads will work on the same row + // - the ones with the same `quad` + auto otherV = __shfl_xor_sync(0xffffffff, myValue, 1); + myValue = fn(myValue, otherV); + otherV = __shfl_xor_sync(0xffffffff, myValue, 2); + myValue = fn(myValue, otherV); + int lane_in_quad = (lane_id & 3); + return lane_in_quad == 0; + } +}; + +template +struct AccumLambdaIteratorSm70 { + static_assert( + cutlass::platform:: + is_same::value, + "only RowMajor is supported"); + + using Policy = typename T::Policy; + using InstructionShape = typename T::InstructionShape; + using OpDelta = typename T::OpDelta; + using Shape = typename T::Shape; + using Element = accum_t; + + static int const kElementsPerPartial = 4; + using EleShapePerPatial = typename cutlass::platform::conditional< + cutlass::platform::is_same::value, + cutlass::MatrixShape<2, 2>, + cutlass::MatrixShape<1, 4>>::type; + static int const kElementsPerMma = 8; + static int const kAccumulatorPatials = 2; + using QuadShapePerPatialMma = cutlass::MatrixShape<4, 4>; + + static int const kRowIters = Policy::TileIterations::kRow * + Policy::MmaIterations::kRow * EleShapePerPatial::kRow; + + static cutlass::MatrixCoord CUTLASS_DEVICE get_lane_offset( + int8_t lane_id, + int8_t warp_id, + typename T::TensorCoord const& tile_offset) { + int quad = (lane_id >> 2); + int lane_in_quad = (lane_id & 3); + int accum_m, accum_n; + + if (cutlass::platform::is_same::value) { + // (quad[2],quad[0])+lane_in_quad[0] + accum_m = (((quad & 0x4) >> 1) + (quad & 0x1)) * 8 + (lane_in_quad & 1); + // (quad[1])+lane_in_quad[1] + accum_n = + ((quad >> 1) & 0x1) * kElementsPerPartial * kAccumulatorPatials + + (lane_in_quad & 2); + } else { + accum_m = (((quad & 0x4) >> 1) + (quad & 0x1)) * 8 + + lane_in_quad; // (quad[2],quad[0]) + accum_n = ((quad >> 1) & 0x1) * kElementsPerPartial * kAccumulatorPatials; + } + return cutlass::MatrixCoord( + accum_m + tile_offset.row() * Shape::kRow, + accum_n + tile_offset.column() * Shape::kColumn); + } + + template + CUTLASS_DEVICE static bool reduceSameRow(int lane_id, DT& myValue, F fn) { + static_assert( + cutlass::platform::is_same::value, + "update to support non-float accum"); + // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-fragment-mma-884-f16 + // T0 & T2 share same line within a quad + auto otherV = __shfl_xor_sync(0xffffffff, myValue, 1 << 1); + myValue = fn(myValue, otherV); + // quad 0 and quad 2 are on the same lines + otherV = __shfl_xor_sync(0xffffffff, myValue, 1 << 3); + myValue = fn(myValue, otherV); + return (lane_id & ((1 << 1) | (1 << 3))) == 0; + } + + template + CUTLASS_DEVICE static void iterateRows( + cutlass::MatrixCoord& lane_offset, + FA beginRow, + FB op, + FC endRow) { + CUTLASS_PRAGMA_UNROLL + for (int tile_m = 0; tile_m < Policy::TileIterations::kRow; ++tile_m) { + CUTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + CUTLASS_PRAGMA_UNROLL + for (int m = 0; m < EleShapePerPatial::kRow; ++m) { + int accum_m = tile_m * Policy::InterleavedTile::kRow + + mma_m * QuadShapePerPatialMma::kRow + m * 2 + lane_offset.row(); + beginRow(accum_m); + + CUTLASS_PRAGMA_UNROLL + for (int tile_n = 0; tile_n < Policy::TileIterations::kColumn; + ++tile_n) { + CUTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; + ++mma_n) { + CUTLASS_PRAGMA_UNROLL + for (int p = 0; p < kAccumulatorPatials; ++p) { + CUTLASS_PRAGMA_UNROLL + for (int n = 0; n < EleShapePerPatial::kColumn; ++n) { + int mma_accum_start = + (((tile_n * Policy::TileIterations::kRow + tile_m) * + Policy::MmaIterations::kColumn + + mma_n) * + Policy::MmaIterations::kRow + + mma_m) * + kElementsPerMma; + int accum_n = tile_n * Policy::InterleavedTile::kColumn + + mma_n * QuadShapePerPatialMma::kColumn + + p * Policy::InterleavedTile::kColumn / 2 + n + + lane_offset.column(); + int idx = mma_accum_start + p * kElementsPerPartial + + m * EleShapePerPatial::kColumn + n; + op(accum_m, accum_n, idx); + } + } + } + } + endRow(accum_m); + } + } + } + } +}; + +template +struct AccumLambdaIteratorSimt { + using Policy = typename T::Policy; + using Iterations = typename T::Iterations; + using Element = typename T::Element; + using Delta = typename T::Delta; + using Shape = typename T::Shape; + static_assert( + cutlass::platform:: + is_same::value, + "only RowMajor is supported"); + + static int const kRowIters = Iterations::kRow * Policy::LaneMmaShape::kM; + + template + CUTLASS_DEVICE static bool reduceSameRow(int lane_id, DT& myValue, F fn) { + CUTLASS_PRAGMA_UNROLL + for (int bit = 1; bit < Policy::WarpShape::kColumn; bit *= 2) { + auto otherV = __shfl_xor_sync(0xffffffff, myValue, bit); + myValue = fn(myValue, otherV); + } + return (lane_id & (Policy::WarpShape::kColumn - 1)) == 0; + } + + template + CUTLASS_DEVICE static void iterateRows( + cutlass::MatrixCoord& lane_offset, + FA beginRow, + FB op, + FC endRow) { + CUTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Iterations::kRow; ++mma_m) { + CUTLASS_PRAGMA_UNROLL + for (int m = 0; m < Policy::LaneMmaShape::kM; ++m) { + int accum_m = mma_m * Delta::kRow + m + lane_offset.row(); + beginRow(accum_m); + + CUTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Iterations::kColumn; ++mma_n) { + int accum_n = + mma_n * Policy::WarpShape::kColumn * Policy::LaneMmaShape::kN + + lane_offset.column(); + CUTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::LaneMmaShape::kN; ++n) { + int idx = n + + Policy::LaneMmaShape::kN * + (mma_n + + Iterations::kColumn * + (m + mma_m * Policy::LaneMmaShape::kM)); + op(accum_m, accum_n + n, idx); + } + } + endRow(accum_m); + } + } + } + + static cutlass::MatrixCoord CUTLASS_DEVICE get_lane_offset( + int8_t lane_id, + int8_t warp_id, + typename T::TensorCoord const& tile_offset) { + static_assert( + cutlass::platform::is_same< + typename Policy::LaneLayout, + cutlass::layout::RowMajorInterleaved<1>>::value, + ""); + typename Policy::LaneLayout lane_layout = Policy::get_lane_layout(); + + cutlass::MatrixCoord lane_offset = lane_layout.inverse(lane_id) * + cutlass::MatrixCoord(Policy::LaneMmaShape::kM, + Policy::LaneMmaShape::kN); + return lane_offset + + tile_offset * cutlass::MatrixCoord(Shape::kRow, Shape::kColumn); + } +}; + +template +struct DefaultMmaAccumLambdaIterator; + +// Simt +template +struct DefaultMmaAccumLambdaIterator< + cutlass::gemm::warp::MmaSimtTileIterator< + S, + cutlass::gemm::Operand::kC, + accum_t, + cutlass::layout::RowMajor, + P, + 1, + 1>, + accum_t, + kWarpSize> { + using WarpIterator = typename cutlass::gemm::warp::MmaSimtTileIterator< + S, + cutlass::gemm::Operand::kC, + accum_t, + cutlass::layout::RowMajor, + P, + 1, + 1>; + using Iterator = AccumLambdaIteratorSimt; +}; + +// TensorOp - Volta +template +struct DefaultMmaAccumLambdaIterator< + cutlass::gemm::warp::MmaVoltaTensorOpAccumulatorTileIterator< + S1, + accum_t, + cutlass::layout::RowMajor, + S2, + cutlass::MatrixShape<1, 1>>, + accum_t, + kWarpSize> { + using WarpIterator = + typename cutlass::gemm::warp::MmaVoltaTensorOpAccumulatorTileIterator< + S1, + accum_t, + cutlass::layout::RowMajor, + S2, + cutlass::MatrixShape<1, 1>>; + using Iterator = AccumLambdaIteratorSm70; +}; + +// TensorOp - Sm75+ +template < + typename S1, + typename S2, + typename S3, + typename accum_t, + int kWarpSize> +struct DefaultMmaAccumLambdaIterator< + cutlass::gemm::warp::MmaTensorOpAccumulatorTileIterator< + S1, + accum_t, + cutlass::layout::RowMajor, + S2, + S3>, + accum_t, + kWarpSize> { + using WarpIterator = + typename cutlass::gemm::warp::MmaTensorOpAccumulatorTileIterator< + S1, + accum_t, + cutlass::layout::RowMajor, + S2, + S3>; + using Iterator = AccumLambdaIteratorSm80; +}; diff --git a/natten/csrc/include/natten/cuda/fna/gemm/mma_from_smem.h b/natten/csrc/include/natten/cuda/fna/gemm/mma_from_smem.h new file mode 100644 index 00000000..9cd71850 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna/gemm/mma_from_smem.h @@ -0,0 +1,1939 @@ +/* + * Copied from xFormers (https://github.com/facebookresearch/xformers/) and + * edited. + * + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + */ +/*************************************************************************************************** + * Copyright (c) 2017 - 2026 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 Tools and utils to store a GEMM output in shmem, and to use that + output as operandA for another GEMM back-to-back +*/ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace cutlass { +namespace gemm { +namespace threadblock { + +/// Shared storage object needed by accumulator +/// From 13_two_tensor_op_fusion/threadblock/b2b_mma_base_smem_accumulator.h +template < + typename Shape_, + typename Element_, + typename Layout_, + typename Padding_> +class AccumulatorSharedStorage { + public: + // + // Type definitions + // + using Shape = Shape_; + using Element = Element_; + using Layout = Layout_; + using Padding = Padding_; + + /// Tensor reference to the accumulator + using TensorRefAccum = cutlass::TensorRef; + + /// Shape of the accumulator matrix in shared memory + using ShapeAccum = cutlass:: + MatrixShape; + + public: + // + // Data members + // + + /// Buffer for accumulator + cutlass::AlignedBuffer accum; + + public: + // + // Methods + // + + /// Returns a layout object for the Accum matrix + CUTLASS_DEVICE + static Layout LayoutAccum() { + return Layout::packed({ShapeAccum::kRow, ShapeAccum::kColumn}); + } + + /// Returns a TensorRef to the Accumulator + CUTLASS_HOST_DEVICE + TensorRefAccum accum_ref() { + return TensorRefAccum{accum.data(), LayoutAccum()}; + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// Taken from +// https://github.com/NVIDIA/cutlass/blob/master/examples/13_two_tensor_op_fusion/threadblock/b2b_mma_base_smem_accumulator.h +//////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + // Maximum K dimension - also the dimension of the shared-memory + // holding `OperandA` + int kMaxK_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Layout in shared-memory of operand A + typename SmemLayoutA, + /// Used for partial specialization + typename Enable = bool> +class MmaBaseFromSharedMemory { + public: + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + static constexpr int kMaxK = kMaxK_; + + ///< Policy describing tuning details + using Policy = Policy_; + + // + // Dependent types + // + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Shape describing the overall GEMM computed from shared memory + /// by each warp. + using WarpGemm = typename Policy::Operator::Shape; + + /// Shape describing the number of warps filling the CTA + using WarpCount = GemmShape< + Shape::kM / WarpGemm::kM, + Shape::kN / WarpGemm::kN, + Shape::kK / WarpGemm::kK>; + using WarpCount1 = WarpCount; + + /// Number of warp-level GEMM oeprations + static int const kWarpGemmIterations = + (WarpGemm::kK / Operator::Policy::MmaShape::kK); + static int const kWarpGemmIterations1 = kWarpGemmIterations; + + /// Number of stages + static int const kStages = Stages; + + /// If this is true, we fill the entire shmem buffer at start + /// and don't need to iterate through it in a circular fashion + static bool const kSmemContainsEntireB = kMaxK <= Shape::kK * kStages; + + /// Tensor reference to the A operand + using TensorRefA = TensorRef; + + /// Tensor reference to the B operand + using TensorRefB = + TensorRef; + + // + // Nested structs + // + + /// Shared storage object needed by threadblock-scoped GEMM + class SharedStorage { + public: + // + // Type definitions + // + + /// Shape of the B matrix operand in shared memory + using ShapeB = MatrixShape< + Shape::kK * kStages + Policy::SmemPaddingB::kRow, + Shape::kN + Policy::SmemPaddingB::kColumn>; + + public: + // + // Data members + // + + /// Buffer for B operand + AlignedBuffer operand_B; + + public: + // + // Methods + // + + /// Returns a layout object for the B matrix + CUTLASS_HOST_DEVICE + static typename Operator::LayoutB LayoutB() { + return Operator::LayoutB::packed({ShapeB::kRow, ShapeB::kColumn}); + } + + /// Returns a TensorRef to the B operand + CUTLASS_HOST_DEVICE + TensorRefB operand_B_ref() { + return TensorRefB{operand_B.data(), LayoutB()}; + } + }; + + protected: + // + // Data members + // + + // /// Iterator to load a warp-scoped tile of A operand from shared memory + // typename Operator::IteratorA warp_tile_iterator_A_; + + /// Iterator to load a warp-scoped tile of B operand from shared memory + typename Operator::IteratorB warp_tile_iterator_B_; + + public: + /// Construct from tensor references + CUTLASS_DEVICE + MmaBaseFromSharedMemory( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + TensorRefB& b_tile, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx) + : warp_tile_iterator_B_(b_tile, lane_idx) {} +}; + +namespace { + +// has necessary trait compliance with WarpIteratorFromSmem but doesn't do +// anything, can be default initialized, and uses fragment that takes up +// (almost) no space. this warp iterator is selected at compile time when +// elementwise on-the-fly scaling for operand A is disabled, in which case +// operations related to loading scale factors for operand A get wiped out by +// the compiler. +template +class NoOpWarpIteratorScale { + public: + // in pipelined+multistage MMA implementations we keep an array of fragments. + // if we aren't using scaling we don't want to waste registers on fragments + // of scale elements, so ideally this would be sized 0. + // Since arrays of zero-sized objects are not allowed, using size as 1. + // The compiler will most likely wipe it out anyways. + using Fragment = cutlass::Array; + + CUTLASS_HOST_DEVICE + NoOpWarpIteratorScale() {} + + CUTLASS_HOST_DEVICE + NoOpWarpIteratorScale(TensorRef const&, int) {} + + CUTLASS_HOST_DEVICE + NoOpWarpIteratorScale& add_tile_offset( + typename TensorRef::TensorCoord const&) { + return *this; + } + + CUTLASS_HOST_DEVICE + NoOpWarpIteratorScale& operator++() { + return *this; + } + + CUTLASS_DEVICE + void load(Fragment&) const {} +}; + +// if scaling is enabled, performs fragment elementwise multiplication between +// fragment and its scaling factor. +template +class FragmentElementwiseScaler; + +// specialization for scaling being enabled. +template +class FragmentElementwiseScaler { + public: + // cast scale_frag to correct type then apply elementwise to fragment + CUTLASS_DEVICE + static Fragment apply(Fragment frag, FragmentScale const& scale_frag) { + Fragment converted_scale_frag = cutlass::NumericArrayConverter< + typename Fragment::Element, + typename FragmentScale::Element, + FragmentScale::kElements>()(scale_frag); + return cutlass::multiplies()(frag, converted_scale_frag); + } +}; + +// specialization for scaling being disabled. doesn't do anything and should +// just get wiped out by the compiler. +template +class FragmentElementwiseScaler { + public: + CUTLASS_DEVICE + static Fragment apply(Fragment frag, FragmentScale const&) { + return frag; + } +}; +} // namespace + +//////////////////////////////////////////////////////////////////////////////// +// Taken from +// https://github.com/NVIDIA/cutlass/blob/master/examples/13_two_tensor_op_fusion/threadblock/b2b_mma_pipelined_smem_accumulator.h +//////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + // BEGIN smem + /// Iterates over the intermediate accumulator tile in shared memory + typename WarpIteratorA_, + /// whether or not to perform elementwise multiplication of A + // by another matrix (A_scale) that is also kept in shared memory prior + // to matmul A @ B + bool ScaleOperandA_, + /// Max GEMM problem size in K dimension + int MaxK, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Transformation applied to B operand + typename TransformB_ = NumericArrayConverter< + typename SmemIteratorB_::Element, + typename IteratorB_::Element, + IteratorB_::Fragment::kElements>, + /// Used for partial specialization + typename Enable = bool> +class MmaPipelinedFromSharedMemory : public MmaBaseFromSharedMemory< + Shape_, + MaxK, + Policy_, + 2, + typename WarpIteratorA_::Layout> { + public: + ///< Base class + using Base = MmaBaseFromSharedMemory< + Shape_, + MaxK, + Policy_, + 2, + typename WarpIteratorA_::Layout>; + + using Shape = + Shape_; ///< Size of the Gemm problem - concept: gemm::GemmShape<> + static constexpr bool ScaleOperandA = ScaleOperandA_; + + using WarpIteratorA = WarpIteratorA_; + ///< loads fragments of A_scale from shared memory if operand A scaling is + ///< enabled. otherwise no-op. + using WarpIteratorAScale = typename cutlass::platform::conditional< + ScaleOperandA, + WarpIteratorA, + NoOpWarpIteratorScale>::type; + + using IteratorB = + IteratorB_; ///< Iterates over tiles of B operand in global memory + using ElementC = ElementC_; ///< Data type of accumulator matrix + using LayoutC = LayoutC_; ///< Layout of accumulator matrix + using Policy = Policy_; ///< Policy describing tuning details + + using SmemIteratorB = SmemIteratorB_; + + using TransformB = TransformB_; + + // + // Dependent types + // + + /// Fragment of operand B loaded from global memory + using FragmentB = typename IteratorB::Fragment; + + /// Fragment of accumulator tile + using FragmentC = typename Policy::Operator::FragmentC; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Obtain the arch tag from the warp-level operator + using ArchTag = typename Policy::Operator::ArchTag; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = Operator::kTransformB; + + // staticaly assert kStages for MmaPipelined is two (Double-buffered pipeline) + static_assert( + (Base::kStages == 2), + "MmaPipelined requires kStages set to value 2"); + + private: + using WarpFragmentA = typename Operator::FragmentA; + + /// fragment type of OperandA elementwise scaling matrix. (almost) empty + /// if operand A scaling is disabled. + using WarpFragmentAScale = typename WarpIteratorAScale::Fragment; + + using WarpFragmentB = typename Operator::FragmentB; + + /// applies scaling factor to operand A fragment if operand A scaling is + /// enabled. otherwise no-op. + using FragmentAScaler = FragmentElementwiseScaler< + WarpFragmentA, + WarpFragmentAScale, + ScaleOperandA>; + + protected: + // /// Iterator to write threadblock-scoped tile of A operand to shared memory + // SmemIteratorA smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB smem_iterator_B_; + + /// Iterator to load a warp-scoped tile of A operand from intermediate + /// accumulator tile + WarpIteratorA warp_tile_iterator_A_; + + /// Iterator to load a warp-scoped tile of A_scale from intermediate + /// accumulator tile (only used if ScaleOperandA_ is true) + WarpIteratorAScale warp_tile_iterator_A_scale_; + + public: + /// constructor for MMA with operand A scaling enabled. + CUTLASS_DEVICE + MmaPipelinedFromSharedMemory( + typename Base::TensorRefA a, // Operand A in shared memory + typename Base::TensorRefA a_scale, // Operand A_scale in shared memory + typename Base::TensorRefB + b_staging, // staging memory for loading tiles of B + int thread_idx, + int warp_idx, + int lane_idx) + : Base(b_staging, thread_idx, warp_idx, lane_idx), + warp_tile_iterator_A_(a, lane_idx), + warp_tile_iterator_A_scale_(a_scale, lane_idx), + smem_iterator_B_(b_staging, thread_idx) { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; + int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset( + {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_A_scale_.add_tile_offset( + {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); + } + + /// Construct from tensor references + CUTLASS_DEVICE + MmaPipelinedFromSharedMemory( + typename Base::TensorRefA a, ///< Operand A in shared memory + typename Base::TensorRefB b_staging, ///< staging memory for loading B + int thread_idx, ///< ID within the threadblock + int warp_idx, ///< ID of warp + int lane_idx) ///< ID of each thread within a warp + : Base(b_staging, thread_idx, warp_idx, lane_idx), + warp_tile_iterator_A_(a, lane_idx), + smem_iterator_B_(b_staging, thread_idx) { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + + int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; + int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset( + {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); + } + + // For API compatibility with MmaMultistageFromSharedMemory + // but not supported as it worsens perf: older gpus < sm80 don't + // support async tranfers and have to waste registers + CUTLASS_DEVICE + void set_prologue_done(bool value) {} + CUTLASS_DEVICE + static void prologue( + typename Base::SharedStorage& shared_storage, + IteratorB iterator_B1, + int thread_idx, + int problem_size_0_n) {} + + CUTLASS_DEVICE + static void drain_cp_asyncs() {} + + /// Perform a threadblock-scoped matrix multiply-accumulate + CUTLASS_DEVICE + void operator()( + int gemm_k_iterations, ///< number of iterations of the mainloop + FragmentC& accum, ///< destination accumulator tile + // IteratorA iterator_A, ///< iterator over A + // operand in global memory + IteratorB iterator_B, ///< iterator over B operand in global memory + FragmentC const& src_accum, ///< source accumulator tile + // TransformA transform_A = TransformA(), ///< transformation + // applied to A fragment + TransformB transform_B = + TransformB()) { ///< transformation applied to B fragment + + // + // Prologue + // + + // Perform accumulation in the 'd' output operand + accum = src_accum; + + FragmentB tb_frag_B; + + tb_frag_B.clear(); + + // The last kblock is loaded in the prolog + iterator_B.set_residual_tile(gemm_k_iterations == 1); + iterator_B.load(tb_frag_B); + + ++iterator_B; + + this->smem_iterator_B_.store(transform_B(tb_frag_B)); + + ++this->smem_iterator_B_; + + __syncthreads(); + + // remember that WarpFragmentAScale and WarpIteratorAScale are empty/no-op + // if scaling is disabled. + + // Pair of fragments used to overlap shared memory loads and math + // instructions + WarpFragmentA warp_frag_A[2]; + WarpFragmentAScale warp_frag_A_scale[2]; + WarpFragmentB warp_frag_B[2]; + warp_frag_A[0].clear(); + warp_frag_A_scale[0].clear(); + warp_frag_B[0].clear(); + + this->warp_tile_iterator_B_.set_kgroup_index(0); + + this->warp_tile_iterator_A_.load(warp_frag_A[0]); + this->warp_tile_iterator_A_scale_.load(warp_frag_A_scale[0]); + this->warp_tile_iterator_B_.load(warp_frag_B[0]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_A_scale_; + ++this->warp_tile_iterator_B_; + + Operator warp_mma; + + int smem_write_stage_idx = 1; + + // Avoid reading out of bounds + iterator_B.set_residual_tile(gemm_k_iterations == 2); + iterator_B.clear_mask(gemm_k_iterations <= 1); + + // Issue loads during the first warp-level matrix multiply-add *AFTER* + // issuing shared memory loads (which have the tightest latency + // requirement). + + // + // Mainloop + // + + // Note: The main loop does not support Base::kWarpGemmIterations == 2. + CUTLASS_GEMM_LOOP + for (; gemm_k_iterations > 0; --gemm_k_iterations) { + // + // Loop over GEMM K dimension + // + + CUTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; + ++warp_mma_k) { + // Load warp-level tiles from shared memory, wrapping to k offset if + // this is the last group as the case may be. + bool hasNext = true; + + if (warp_mma_k == Base::kWarpGemmIterations - 1) { + if (gemm_k_iterations > 1) { + // Write fragments to shared memory + this->smem_iterator_B_.store(transform_B(tb_frag_B)); + } + + __syncthreads(); + + ++this->smem_iterator_B_; + + // Add negative offsets to return iterators to the 'start' of the + // circular buffer in shared memory SMEM: Don't reset iterator A, as + // we are continuing our iteration at this point + if (smem_write_stage_idx == 1) { + this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); + } else { + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations, + 0}); + } + + smem_write_stage_idx ^= 1; + hasNext = gemm_k_iterations > 1; + } + + // Only read the next if we need to + if (hasNext) { + this->warp_tile_iterator_B_.set_kgroup_index( + (warp_mma_k + 1) % Base::kWarpGemmIterations); + + this->warp_tile_iterator_A_.load(warp_frag_A[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_A_scale_.load( + warp_frag_A_scale[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_B_.load(warp_frag_B[(warp_mma_k + 1) % 2]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_A_scale_; + ++this->warp_tile_iterator_B_; + + if (warp_mma_k == 0) { + iterator_B.load(tb_frag_B); + + ++iterator_B; + + // Avoid reading out of bounds if this was the last loop iteration + iterator_B.set_residual_tile(gemm_k_iterations == 3); + iterator_B.clear_mask(gemm_k_iterations <= 2); + } + } + + warp_mma( + accum, + FragmentAScaler::apply( + warp_frag_A[warp_mma_k % 2], warp_frag_A_scale[warp_mma_k % 2]), + warp_frag_B[warp_mma_k % 2], + accum); + } + } + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// Taken from +// https://github.com/NVIDIA/cutlass/blob/master/examples/13_two_tensor_op_fusion/threadblock/b2b_mma_multistage_smem_accumulator.h +//////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape1_, + /// Iterates over the intermediate accumulator tile in shared memory + typename WarpIteratorA1_, + /// whether or not to perform elementwise multiplication of A + // by another matrix (A_scale) that is also kept in shared memory prior + // to matmul A @ B + bool ScaleOperandA_, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB1_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB1_, + /// Cache operation for operand B + cutlass::arch::CacheOperation::Kind CacheOpB1, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy1_, + /// Number of stages, + int Stages_, + int kMaxK_, + /// Used for partial specialization + typename Enable = bool> +class MmaMultistageFromSharedMemory : public MmaBaseFromSharedMemory< + Shape1_, + kMaxK_, + Policy1_, + Stages_, + typename WarpIteratorA1_::Layout> { + public: + ///< Base class + using Base = MmaBaseFromSharedMemory< + Shape1_, + kMaxK_, + Policy1_, + Stages_, + typename WarpIteratorA1_::Layout>; + + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape1 = Shape1_; + ///< Iterates over tiles of B operand in global memory + using IteratorB1 = IteratorB1_; + using IteratorB = IteratorB1; + ///< Policy describing tuning details + using Policy1 = Policy1_; + + using SmemIteratorB1 = SmemIteratorB1_; + using WarpIteratorA1 = WarpIteratorA1_; ///< Iterates over the intermediate + ///< accumulator tile in shared memory + static constexpr bool ScaleOperandA = ScaleOperandA_; + + ///< warp level iterator over A_scale matrix tile kept in shared memory. + ///< if elementwise A scaling is disabled then everything this does is no-op. + using WarpIteratorAScale = typename cutlass::platform::conditional< + ScaleOperandA, + WarpIteratorA1, + NoOpWarpIteratorScale>::type; + ///< Data type of accumulator matrix + using ElementC = ElementC_; + ///< Layout of accumulator matrix + using LayoutC = LayoutC_; + + static cutlass::arch::CacheOperation::Kind const kCacheOpB1 = CacheOpB1; + static constexpr bool kSmemContainsEntireB = Base::kSmemContainsEntireB; + + // + // Dependent types + // + + /// Fragment of accumulator tile + using FragmentC1 = typename Policy1::Operator::FragmentC; + using FragmentC = FragmentC1; + + /// Warp-level Mma + using Operator1 = typename Policy1::Operator; + + /// Minimum architecture is Sm80 to support cp.async + using ArchTag = arch::Sm80; + + /// Complex transform on B operand + static ComplexTransform const kTransformB1 = Operator1::kTransformB; + + /// Internal structure exposed for introspection. + struct Detail { + static_assert( + Base::kWarpGemmIterations1 > 1, + "The pipelined structure requires at least two warp-level " + "GEMM operations."); + + /// Number of cp.async instructions to load one stage of operand B + static int const TBLoadIterationsB1 = + IteratorB1::ThreadMap::Iterations::kCount; + + /// Number of cp.async instructions to load on group of operand B + static int const kAccessesPerGroupB1 = + (TBLoadIterationsB1 + Base::kWarpGemmIterations1 - 1) / + Base::kWarpGemmIterations1; + }; + + static constexpr int kNumStagesConcurrentLoad = + kSmemContainsEntireB ? Base::kStages : Base::kStages - 1; + + private: + using WarpLoadedFragmentA1 = typename Operator1::FragmentA; + /// fragment of OperandA scale matrix. if operand A scaling is disabled this + /// is (almost) empty. + using WarpLoadedFragmentA1Scale = typename WarpIteratorAScale::Fragment; + using WarpLoadedFragmentB1 = typename Operator1::FragmentB; + using WarpTransformedFragmentA1 = typename Operator1::TransformedFragmentA; + using WarpTransformedFragmentB1 = typename Operator1::TransformedFragmentB; + + /// applies elementwise scaling to fragment of A. if operand A scaling is + /// disabled this is a no-op. + using FragmentAScaler = FragmentElementwiseScaler< + WarpLoadedFragmentA1, + WarpLoadedFragmentA1Scale, + ScaleOperandA>; + + private: + // + // Data members + // + + /// Iterator to load a warp-scoped tile of A1 operand from intermediate + /// accumulator tile + WarpIteratorA1 warp_tile_iterator_A1_; + + /// Iterator to load a warp-scoped tile of A1_scale operand from shared memory + /// if operand A scaling is disabled everything this does is a no-op. + WarpIteratorAScale warp_tile_iterator_A1_scale_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB1 smem_iterator_B1_; + + bool prologue_done_; + + public: + /// constructor for MMA with operand A scaling enabled. + CUTLASS_DEVICE + MmaMultistageFromSharedMemory( + typename Base::TensorRefA a, + typename Base::TensorRefA a_scale, + typename Base::TensorRefB b_tile, + int thread_idx, + int warp_idx, + int lane_idx) + : Base(b_tile, thread_idx, warp_idx, lane_idx), + warp_tile_iterator_A1_(a, lane_idx), + warp_tile_iterator_A1_scale_(a_scale, lane_idx), + smem_iterator_B1_(b_tile, thread_idx), + prologue_done_(false) { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + int warp_idx_mn_1 = + warp_idx % (Base::WarpCount1::kM * Base::WarpCount1::kN); + int warp_idx_k_1 = warp_idx / (Base::WarpCount1::kM * Base::WarpCount1::kN); + int warp_idx_m_1 = warp_idx_mn_1 % Base::WarpCount1::kM; + int warp_idx_n_1 = warp_idx_mn_1 / Base::WarpCount1::kM; + + // Add per-warp offsets in units of warp-level tiles + warp_tile_iterator_A1_.add_tile_offset( + {warp_idx_m_1, Base::kWarpGemmIterations1 * warp_idx_k_1}); + warp_tile_iterator_A1_scale_.add_tile_offset( + {warp_idx_m_1, Base::kWarpGemmIterations1 * warp_idx_k_1}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations1 * warp_idx_k_1, warp_idx_n_1}); + } + + /// Construct from tensor references + CUTLASS_DEVICE + MmaMultistageFromSharedMemory( + typename Base::TensorRefA a, + typename Base::TensorRefB b_tile, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx) + : Base(b_tile, thread_idx, warp_idx, lane_idx), + warp_tile_iterator_A1_(a, lane_idx), + smem_iterator_B1_(b_tile, thread_idx), + prologue_done_(false) { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn_1 = + warp_idx % (Base::WarpCount1::kM * Base::WarpCount1::kN); + int warp_idx_k_1 = warp_idx / (Base::WarpCount1::kM * Base::WarpCount1::kN); + + int warp_idx_m_1 = warp_idx_mn_1 % Base::WarpCount1::kM; + int warp_idx_n_1 = warp_idx_mn_1 / Base::WarpCount1::kM; + + // Add per-warp offsets in units of warp-level tiles + warp_tile_iterator_A1_.add_tile_offset( + {warp_idx_m_1, Base::kWarpGemmIterations1 * warp_idx_k_1}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations1 * warp_idx_k_1, warp_idx_n_1}); + } + + CUTLASS_DEVICE + void set_prologue_done(bool value) { + prologue_done_ = value; + } + + CUTLASS_DEVICE + static void prologue( + typename Base::SharedStorage& shared_storage, + IteratorB iterator_B1, + int thread_idx, + int problem_size_0_n) { + SmemIteratorB1 smem_iterator_B1(shared_storage.operand_B_ref(), thread_idx); + _prologue( + iterator_B1, + (problem_size_0_n + Base::Shape::kK - 1) / Base::Shape::kK, + smem_iterator_B1); + } + + CUTLASS_DEVICE + static void drain_cp_asyncs() { + // commit and drain all pending and predicated cp.async pnz from the GEMM + // mainloop + cutlass::arch::cp_async_fence(); + cutlass::arch::cp_async_wait<0>(); + __syncthreads(); + } + + CUTLASS_DEVICE + void copy_tiles_and_advance_1( + IteratorB1& iterator_B1, + int group_start_B1 = 0) { + iterator_B1.set_iteration_index( + group_start_B1 * IteratorB1::kAccessesPerVector); + this->smem_iterator_B1_.set_iteration_index(group_start_B1); + + // Load for operand B + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupB1; ++j) { + if (group_start_B1 + j < Detail::TBLoadIterationsB1) { + typename IteratorB1::AccessType* dst_ptr = + reinterpret_cast( + this->smem_iterator_B1_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorB1::ThreadMap::kElementsPerAccess / + IteratorB1::kAccessesPerVector / 8; + + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB1::kAccessesPerVector; ++v) { + auto gmem_ptr = iterator_B1.get(); + + cutlass::arch::cp_async_zfill( + dst_ptr + v, gmem_ptr, iterator_B1.valid()); + + ++iterator_B1; + } + ++this->smem_iterator_B1_; + } + } + } + + CUTLASS_DEVICE + static void _prologue( + IteratorB& iterator_B1, + int32_t gemm_k_iterations_1, + SmemIteratorB1& smem_iterator_B1_) { + // Issue several complete stages + CUTLASS_PRAGMA_UNROLL + for (int stage = 0; stage < kNumStagesConcurrentLoad; + ++stage, --gemm_k_iterations_1) { + iterator_B1.set_residual_tile(gemm_k_iterations_1 == 1); + iterator_B1.clear_mask(gemm_k_iterations_1 == 0); + + iterator_B1.set_iteration_index(0); + smem_iterator_B1_.set_iteration_index(0); + + // Load for operand B + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::TBLoadIterationsB1; ++j) { + typename IteratorB1::AccessType* dst_ptr = + reinterpret_cast( + smem_iterator_B1_.get()); + + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB1::kAccessesPerVector; ++v) { + int const kSrcBytes = + sizeof_bits::value * + IteratorB1::ThreadMap::kElementsPerAccess / + IteratorB1::kAccessesPerVector / 8; + + cutlass::arch::cp_async_zfill( + dst_ptr + v, iterator_B1.get(), iterator_B1.valid()); + + ++iterator_B1; + } + + ++smem_iterator_B1_; + } + + // Move to the next stage + iterator_B1.add_tile_offset({1, 0}); + + smem_iterator_B1_.add_tile_offset({1, 0}); + + // Defines the boundary of a stage of cp.async. + cutlass::arch::cp_async_fence(); + } + iterator_B1.set_residual_tile(gemm_k_iterations_1 == 1); + iterator_B1.clear_mask(gemm_k_iterations_1 == 0); + } + + /// Perform a threadblock-scoped matrix multiply-accumulate + CUTLASS_DEVICE + void operator()( + ///< problem size of GEMM + int gemm_k_iterations_1_, + ///< destination accumulator tile + FragmentC1& accum, + ///< iterator over B1 operand in global memory + IteratorB1 iterator_B1, + ///< initial value of accumulator + FragmentC1 const& src_accum) { + // 2nd Gemm + + // + // Prologue + // + // Perform accumulation in the 'd' output operand + accum = src_accum; + + if (!prologue_done_) { + _prologue(iterator_B1, gemm_k_iterations_1_, smem_iterator_B1_); + } else if (!kSmemContainsEntireB) { + // Restore the iterators increments + + int gemm_k_iterations_1 = gemm_k_iterations_1_; + // Issue several complete stages + CUTLASS_PRAGMA_UNROLL + for (int stage = 0; stage < kNumStagesConcurrentLoad; + ++stage, --gemm_k_iterations_1) { + iterator_B1.set_iteration_index(0); + this->smem_iterator_B1_.set_iteration_index(0); + + // Load for operand B + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::TBLoadIterationsB1; ++j) { + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB1::kAccessesPerVector; ++v) { + ++iterator_B1; + } + ++this->smem_iterator_B1_; + } + iterator_B1.add_tile_offset({1, 0}); + this->smem_iterator_B1_.add_tile_offset({1, 0}); + } + iterator_B1.set_residual_tile(gemm_k_iterations_1 <= 1); + iterator_B1.clear_mask(gemm_k_iterations_1 <= 0); + } + + // DEPBAR+SYNC + cutlass::arch::cp_async_wait(); + __syncthreads(); + + // remember that WarpFragmentAScale and WarpIteratorAScale are no-op/empty + // if scaling is disabled. + + // Pair of fragments used to overlap shared memory loads and math + // instructions + WarpLoadedFragmentA1 warp_loaded_frag_A1[2]; + WarpLoadedFragmentA1Scale warp_loaded_frag_A1_scale[2]; + WarpLoadedFragmentB1 warp_loaded_frag_B1[2]; + WarpTransformedFragmentA1 warp_transformed_frag_A1[2]; + WarpTransformedFragmentB1 warp_transformed_frag_B1[2]; + + Operator1 warp_mma1; + + warp_tile_iterator_A1_.load(warp_loaded_frag_A1[0]); + ++warp_tile_iterator_A1_; + + warp_tile_iterator_A1_scale_.load(warp_loaded_frag_A1_scale[0]); + ++warp_tile_iterator_A1_scale_; + + this->warp_tile_iterator_B_.set_kgroup_index(0); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B1[0]); + ++this->warp_tile_iterator_B_; + + int smem_write_stage_idx = Base::kStages - 1; + int smem_read_stage_idx = 0; + + warp_mma1.transform( + warp_transformed_frag_A1[0], + warp_transformed_frag_B1[0], + FragmentAScaler::apply( + warp_loaded_frag_A1[0], warp_loaded_frag_A1_scale[0]), + warp_loaded_frag_B1[0]); + + // tf32x3 kernels use staging accumulation. warp_mma uses a temporary + // accumulator and this temporary accumulator is added to the final + // accumulator once in every mainloop iteration. + plus plus_accum; + + FragmentC1 tmp_accum; + + if (platform::is_same< + typename Operator1::MathOperator, + arch::OpMultiplyAddFastF32>::value || + platform::is_same< + typename Operator1::MathOperator, + arch::OpMultiplyAddComplexFastF32>::value) { + tmp_accum.clear(); + } + + // + // Mainloop + // + + CUTLASS_PRAGMA_UNROLL + for (int gemm_k_iterations_1 = gemm_k_iterations_1_ - (Base::kStages - 1); + gemm_k_iterations_1 > (-Base::kStages + 1); + gemm_k_iterations_1--) { + // + // Loop over GEMM K dimension + // + + // Computes a warp-level GEMM on data held in shared memory + // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate + CUTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations1; + ++warp_mma_k) { + // Load warp-level tile from accumulator fragment (A) + // or shared memory (operand B) + this->warp_tile_iterator_B_.set_kgroup_index( + (warp_mma_k + 1) % Base::kWarpGemmIterations1); + // skip warp tile loading for the last kgroup (we are out of the buf) + if (gemm_k_iterations_1 > (-Base::kStages + 2) || + warp_mma_k < Base::kWarpGemmIterations1 - 1) { + warp_tile_iterator_A1_.load( + warp_loaded_frag_A1[(warp_mma_k + 1) % 2]); + warp_tile_iterator_A1_scale_.load( + warp_loaded_frag_A1_scale[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_B_.load( + warp_loaded_frag_B1[(warp_mma_k + 1) % 2]); + } + ++warp_tile_iterator_A1_; + ++warp_tile_iterator_A1_scale_; + ++this->warp_tile_iterator_B_; + + if (warp_mma_k > 0) + warp_mma1.transform( + warp_transformed_frag_A1[warp_mma_k % 2], + warp_transformed_frag_B1[warp_mma_k % 2], + FragmentAScaler::apply( + warp_loaded_frag_A1[warp_mma_k % 2], + warp_loaded_frag_A1_scale[warp_mma_k % 2]), + warp_loaded_frag_B1[warp_mma_k % 2]); + + if (platform::is_same< + typename Operator1::MathOperator, + arch::OpMultiplyAddFastF32>::value || + platform::is_same< + typename Operator1::MathOperator, + arch::OpMultiplyAddComplexFastF32>::value) { + warp_mma1( + tmp_accum, + warp_transformed_frag_A1[warp_mma_k % 2], + warp_transformed_frag_B1[warp_mma_k % 2], + tmp_accum); + + if (warp_mma_k == 0) { + accum = plus_accum(accum, tmp_accum); + tmp_accum.clear(); + } + } else { + warp_mma1( + accum, + warp_transformed_frag_A1[warp_mma_k % 2], + warp_transformed_frag_B1[warp_mma_k % 2], + accum); + } + + // Issue global->shared copies for the this stage + if (warp_mma_k < Base::kWarpGemmIterations1 - 1) { + int group_start_iteration_B1; + + group_start_iteration_B1 = warp_mma_k * Detail::kAccessesPerGroupB1; + + if (!kSmemContainsEntireB) { + copy_tiles_and_advance_1(iterator_B1, group_start_iteration_B1); + } + } + + if (warp_mma_k + 2 == Base::kWarpGemmIterations1) { + int group_start_iteration_B1; + group_start_iteration_B1 = + (warp_mma_k + 1) * Detail::kAccessesPerGroupB1; + + if (!kSmemContainsEntireB) { + copy_tiles_and_advance_1(iterator_B1, group_start_iteration_B1); + } + + // Inserts a memory fence between stages of cp.async instructions. + cutlass::arch::cp_async_fence(); + + // Waits until kStages-2 stages have committed. + arch::cp_async_wait(); + __syncthreads(); + + // Move to the next stage + iterator_B1.add_tile_offset({1, 0}); + + this->smem_iterator_B1_.add_tile_offset({1, 0}); + + // Add negative offsets to return iterators to the 'start' of the + // circular buffer in shared memory + if (!kSmemContainsEntireB) { + if (smem_write_stage_idx == (Base::kStages - 1)) { + this->smem_iterator_B1_.add_tile_offset({-Base::kStages, 0}); + smem_write_stage_idx = 0; + } else { + ++smem_write_stage_idx; + } + + if (smem_read_stage_idx == (Base::kStages - 1)) { + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy1::kPartitionsK * + Base::kWarpGemmIterations1, + 0}); + smem_read_stage_idx = 0; + } else { + ++smem_read_stage_idx; + } + } + + iterator_B1.set_residual_tile(gemm_k_iterations_1 == 2); + iterator_B1.clear_mask(gemm_k_iterations_1 == 1); + } + + // Do any conversions feeding the first stage at the end of the loop so + // we can start right away on mma instructions + if (warp_mma_k + 1 == Base::kWarpGemmIterations1) + warp_mma1.transform( + warp_transformed_frag_A1[(warp_mma_k + 1) % 2], + warp_transformed_frag_B1[(warp_mma_k + 1) % 2], + FragmentAScaler::apply( + warp_loaded_frag_A1[(warp_mma_k + 1) % 2], + warp_loaded_frag_A1_scale[(warp_mma_k + 1) % 2]), + warp_loaded_frag_B1[(warp_mma_k + 1) % 2]); + } + } + + if (platform::is_same< + typename Operator1::MathOperator, + arch::OpMultiplyAddFastF32>::value || + platform::is_same< + typename Operator1::MathOperator, + arch::OpMultiplyAddComplexFastF32>::value) { + accum = plus_accum(accum, tmp_accum); + } + } +}; + +// Converts a "regular" Mma into their counterpart from shared memory +template < + int NADim, + typename Mma_, + int kMaxK, + typename WarpIteratorA_, + /// whether or not to apply elementwise multiplication of operand A by + /// another matrix in shared memory before usage in A @ B + bool kScaleOperandA, + bool kTransposeA = false> +struct DefaultMmaFromSharedMemory; + +// Mma pipelined +template < + int NADim, + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + typename WarpIteratorA_, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Transformation applied to A operand + typename TransformA_, + /// Transformation applied to B operand + typename TransformB_, + // Max MMA problem size K + int kMaxK, + /// whether or not to apply elementwise multiplication of operand A by + /// another matrix in shared memory before usage in A @ B + bool kScaleOperandA, + bool kTransposeA> +struct DefaultMmaFromSharedMemory< + NADim, + MmaPipelined< + Shape_, + IteratorA_, + SmemIteratorA_, + IteratorB_, + SmemIteratorB_, + ElementC_, + LayoutC_, + Policy_, + TransformA_, + TransformB_>, + kMaxK, + WarpIteratorA_, + kScaleOperandA, + kTransposeA> { + using RegularMma = MmaPipelined< + Shape_, + IteratorA_, + SmemIteratorA_, + IteratorB_, + SmemIteratorB_, + ElementC_, + LayoutC_, + Policy_, + TransformA_, + TransformB_>; + + using WarpShape = typename Policy_::Operator::Shape; + using InstructionShape = typename Policy_::Operator::InstructionShape; + using ArchMmaOperator = typename Policy_::Operator; + + static constexpr bool kIsTransposedA = false; + using WarpIteratorA = WarpIteratorA_; + using IteratorB = typename cutlass::transform::threadblock:: + MakeIteratorResidualLast::Iterator; + + using Mma = typename cutlass::gemm::threadblock::MmaPipelinedFromSharedMemory< + Shape_, + WarpIteratorA, + kScaleOperandA, + kMaxK, + IteratorB, + SmemIteratorB_, + ElementC_, + LayoutC_, + Policy_>; +}; + +template < + int NADim, + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + typename WarpIteratorA_, + /// Cache operation for operand A + cutlass::arch::CacheOperation::Kind CacheOpA, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Cache operation for operand B + cutlass::arch::CacheOperation::Kind CacheOpB, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear, + int kMaxK, + /// whether or not to apply elementwise multiplication of operand A by + /// another matrix in shared memory before usage in A @ B + bool kScaleOperandA, + bool kTransposeA> +struct DefaultMmaFromSharedMemory< + NADim, + MmaMultistage< + Shape_, + IteratorA_, + SmemIteratorA_, + CacheOpA, + IteratorB_, + SmemIteratorB_, + CacheOpB, + ElementC_, + LayoutC_, + Policy_, + Stages, + SharedMemoryClear>, + kMaxK, + WarpIteratorA_, + kScaleOperandA, + kTransposeA> { + using RegularMma = MmaMultistage< + Shape_, + IteratorA_, + SmemIteratorA_, + CacheOpA, + IteratorB_, + SmemIteratorB_, + CacheOpB, + ElementC_, + LayoutC_, + Policy_, + Stages, + SharedMemoryClear>; + + using WarpShape = typename Policy_::Operator::Shape; + using InstructionShape = typename Policy_::Operator::InstructionShape; + using WarpIteratorTranspose = TransposeWarpIterator; + static constexpr bool kIsTransposedA = + WarpIteratorTranspose::kSupportsTranspose && kTransposeA; + using WarpIteratorA = typename platform::conditional< + kIsTransposedA, + typename WarpIteratorTranspose::Iterator, + WarpIteratorA_>::type; + + // Reduce the number of stages if we don't need that many + static int constexpr kStagesMax = + (kMaxK + int(Shape_::kK) - 1) / int(Shape_::kK); + static int constexpr kStages = cutlass::const_min(Stages, kStagesMax); + + using IteratorB = typename cutlass::transform::threadblock:: + MakeIteratorResidualLast::Iterator; + using Mma = + typename cutlass::gemm::threadblock::MmaMultistageFromSharedMemory< + Shape_, + WarpIteratorA, + kScaleOperandA, + IteratorB, + SmemIteratorB_, + RegularMma::kCacheOpB, + ElementC_, + LayoutC_, + Policy_, + kStages, + kMaxK>; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + int NADim, + typename IteratorC, + typename Operator, + typename scalar_t, + typename WarpShape_, + typename ThreadblockShape_> +struct B2bGemm; + +// Tensor Cores >= Sm75 specialization (Ampere ...) +template < /// Size of the matrix to load (concept: MatrixShape) + int NADim, + typename Shape_, + /// Element type + typename Element_, + /// Layout of operand in memory + typename Layout_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions, concept: MatrixShape) + typename OpDelta_, + typename Operator, + typename scalar_t, + typename WarpShape_, + typename ThreadblockShape_> +struct B2bGemm< + NADim, + cutlass::gemm::warp::MmaTensorOpAccumulatorTileIterator< + Shape_, + Element_, + Layout_, + InstructionShape_, + OpDelta_>, + Operator, + scalar_t, + WarpShape_, + ThreadblockShape_> { + static_assert(NADim >= 1 && NADim < 4); + using Dim = typename natten::cuda::fna::GetDim::type; + + using IteratorC = + typename cutlass::gemm::warp::MmaTensorOpAccumulatorTileIterator< + Shape_, + Element_, + Layout_, + InstructionShape_, + OpDelta_>; + using FragmentC = typename IteratorC::Fragment; + using InstructionShape = InstructionShape_; + using WarpShape = WarpShape_; + using ThreadblockShape = ThreadblockShape_; + using accum_t = Element_; + using lse_scalar_t = float; + + using SmemAccumulatorLayout = cutlass::layout::RowMajor; + + // Iterator to load accumulators (results of matmul in registers) + using FragmentIteratorAccumulator = + cutlass::epilogue::warp::FragmentIteratorTensorOp< + WarpShape, + InstructionShape, + accum_t, + typename Operator::Policy::Operator::FragmentC, + cutlass::layout::RowMajor>; + + // Iterator to store to shared-memory + using SmemIteratorD0 = typename cutlass::epilogue::warp::TileIteratorTensorOp< + WarpShape, + InstructionShape, + scalar_t, // accum_t, + SmemAccumulatorLayout>; + using AccumulatorSharedStorage = + cutlass::gemm::threadblock::AccumulatorSharedStorage< + ThreadblockShape, + typename SmemIteratorD0::Element, + typename SmemIteratorD0::TensorLayout, + typename SmemIteratorD0::Padding>; + // We need to provide an operation for the epilogue. Let's create an + // operation that does nothing (ScaleType::Nothing), just converts + // from accum_t (float) -> scalar_t (can be half) + using OutputOpNoOp = cutlass::epilogue::thread::LinearCombination< + typename SmemIteratorD0::Element, // ElementOutput + FragmentIteratorAccumulator::Fragment::kElements, + accum_t, // ElementAccumulator + typename SmemIteratorD0::Element, // ElementCompute + cutlass::epilogue::thread::ScaleType::Nothing>; + using Epilogue = cutlass::epilogue::threadblock::EpilogueSmemAccumulator< + SmemIteratorD0, + FragmentIteratorAccumulator, + SmemIteratorD0, // ScaleBiasIterator - not used + OutputOpNoOp>; + + // Epilogue 2: with LSE (for backwards pass) + using EpilogueOpApplyExp = cutlass::epilogue::thread::ApplyExp< + scalar_t, // ElementOutput_ + lse_scalar_t, // ElementLSE_ + accum_t, // ElementAccumulator_ + accum_t, // ElementCompute_ + 128 / cutlass::sizeof_bits::value>; + using EpilogueWithExp = + cutlass::epilogue::threadblock::EpilogueSmemAccumulator< + SmemIteratorD0, // ScaleBiasIterator - not used + FragmentIteratorAccumulator, + SmemIteratorD0, // ScaleBiasIterator - not used + EpilogueOpApplyExp>; + + static void CUTLASS_DEVICE accumToSmem( + AccumulatorSharedStorage& shared_storage, + FragmentC const& accum, + int lane_id, + cutlass::MatrixCoord const& tile_coords) { + SmemIteratorD0 smem_iterator_attn(shared_storage.accum_ref(), lane_id); + smem_iterator_attn.add_tile_offset( + tile_coords * + cutlass::MatrixCoord{ + SmemIteratorD0::TileIterations::kRow, + SmemIteratorD0::TileIterations::kColumn}); + Epilogue epilogue; + epilogue(OutputOpNoOp({}), smem_iterator_attn, accum); + } + + // NOTE(alih): we only apply the exp operator here; + // LSE is pre-fetched into shared memory, and + // subtracted when we check the NA mask and set + // invalid weights to -info. + // Refer to kernel_backward.h for more details. + static void CUTLASS_DEVICE accumApplyExpToSmem( + AccumulatorSharedStorage& shared_storage, + FragmentC& accum, + int thread_id, + int warp_id, + int lane_id, + cutlass::MatrixCoord const& tile_coords) { + SmemIteratorD0 smem_iterator_attn(shared_storage.accum_ref(), lane_id); + smem_iterator_attn.add_tile_offset( + tile_coords * + cutlass::MatrixCoord{ + SmemIteratorD0::TileIterations::kRow, + SmemIteratorD0::TileIterations::kColumn}); + EpilogueWithExp epilogue; + EpilogueOpApplyExp exp({}); + epilogue(exp, smem_iterator_attn, accum); + } +}; + +// Volta Specialization +// only supported for f16 +template < + int NADim, + typename Operator, + typename WarpShape_, + typename ThreadblockShape_> +struct B2bGemm< + NADim, + cutlass::gemm::warp::MmaVoltaTensorOpAccumulatorTileIterator< + cutlass::MatrixShape<32, 32>, + float, + cutlass::layout::RowMajor, + cutlass::gemm::GemmShape<16, 16, 4>, + cutlass::MatrixShape<1, 1>>, + Operator, + cutlass::half_t, + WarpShape_, + ThreadblockShape_> { + static_assert(NADim >= 1 && NADim < 4); + using Dim = typename natten::cuda::fna::GetDim::type; + + using IteratorC = + cutlass::gemm::warp::MmaVoltaTensorOpAccumulatorTileIterator< + cutlass::MatrixShape<32, 32>, + float, + cutlass::layout::RowMajor, + cutlass::gemm::GemmShape<16, 16, 4>, + cutlass::MatrixShape<1, 1>>; + using scalar_t = cutlass::half_t; + using accum_t = IteratorC::Element; + using WarpShape = WarpShape_; + using ThreadblockShape = ThreadblockShape_; + using FragmentC = IteratorC::Fragment; + using lse_scalar_t = float; + + // Storage in shared-memory for Q.Kt + using SmemAccumulatorLayout = + cutlass::layout::RowMajorVoltaTensorOpMultiplicandCrosswise<16, 32>; + using AccumulatorSharedStorage = + cutlass::gemm::threadblock::AccumulatorSharedStorage< + ThreadblockShape, + scalar_t, + SmemAccumulatorLayout, + cutlass::MatrixShape<0, 0> // Padding + >; + using TensorRef = cutlass::TensorRef; + using Policy = typename IteratorC::Policy; + using Element = accum_t; + // Those are MmaVoltaTensorOpAccumulatorTileIterator private fields + // Let's copy their values + static int const kElementsPerPartial = 4; + using EleShapePerPatial = typename cutlass::platform::conditional< + cutlass::platform::is_same::value, + cutlass::MatrixShape<2, 2>, + cutlass::MatrixShape<1, 4>>::type; + static int const kElementsPerMma = 8; + static int const kAccumulatorPatials = 2; + using QuadShapePerPatialMma = cutlass::MatrixShape<4, 4>; + + static void CUTLASS_DEVICE accumToSmem( + AccumulatorSharedStorage& shared_storage, + FragmentC const& accum, + int lane_id, + cutlass::MatrixCoord const& tile_coords) { + // ctor - from MmaVoltaTensorOpAccumulatorTileIterator + TensorRef ref_(shared_storage.accum_ref()); + int quad = (lane_id >> 2); + int lane_in_quad = (lane_id & 3); + int accum_m, accum_n; + + if (cutlass::platform::is_same::value) { + // (quad[2],quad[0])+lane_in_quad[0] + accum_m = (((quad & 0x4) >> 1) + (quad & 0x1)) * 8 + (lane_in_quad & 1); + // (quad[1])+lane_in_quad[1] + accum_n = + ((quad >> 1) & 0x1) * kElementsPerPartial * kAccumulatorPatials + + (lane_in_quad & 2); + } else { + accum_m = (((quad & 0x4) >> 1) + (quad & 0x1)) * 8 + + lane_in_quad; // (quad[2],quad[0]) + accum_n = ((quad >> 1) & 0x1) * kElementsPerPartial * kAccumulatorPatials; + } + cutlass::MatrixCoord lane_offset(accum_m, accum_n); + + // Tile offset + ref_.add_coord_offset( + tile_coords * + cutlass::MatrixCoord( + {IteratorC::Shape::kRow, IteratorC::Shape::kColumn})); + + using AccessType = cutlass::Array; + + // store - from MmaVoltaTensorOpAccumulatorTileIterator + CUTLASS_PRAGMA_UNROLL + for (int tile_n = 0; tile_n < Policy::TileIterations::kColumn; ++tile_n) { + CUTLASS_PRAGMA_UNROLL + for (int tile_m = 0; tile_m < Policy::TileIterations::kRow; ++tile_m) { + CUTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + CUTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + int mma_accum_start = + (((tile_n * Policy::TileIterations::kRow + tile_m) * + Policy::MmaIterations::kColumn + + mma_n) * + Policy::MmaIterations::kRow + + mma_m) * + kElementsPerMma; + + CUTLASS_PRAGMA_UNROLL + for (int p = 0; p < kAccumulatorPatials; ++p) { + CUTLASS_PRAGMA_UNROLL + for (int m = 0; m < EleShapePerPatial::kRow; ++m) { + int accum_m = tile_m * Policy::InterleavedTile::kRow + + mma_m * QuadShapePerPatialMma::kRow + m * 2; + int accum_n = tile_n * Policy::InterleavedTile::kColumn + + mma_n * QuadShapePerPatialMma::kColumn + + p * Policy::InterleavedTile::kColumn / 2; + int r = (accum_m + lane_offset.row()); + AccessType to_store; + CUTLASS_PRAGMA_UNROLL + for (int n = 0; n < EleShapePerPatial::kColumn; ++n) { + int idx = mma_accum_start + p * kElementsPerPartial + + m * EleShapePerPatial::kColumn + n; + int c = (accum_n + n + lane_offset.column()); + to_store[n] = scalar_t(accum[idx]); + } + int c = (accum_n + lane_offset.column()); + assert(r < 32); + assert(c < 32); + *reinterpret_cast( + ref_.data() + ref_.offset({r, c})) = to_store; + } + } + } + } + } + } + } + + // NOTE(alih): we only apply the exp operator here; + // LSE is pre-fetched into shared memory, and + // subtracted when we check the NA mask and set + // invalid weights to -info. + // Refer to kernel_backward.h for more details. + static void CUTLASS_DEVICE accumApplyExpToSmem( + AccumulatorSharedStorage& shared_storage, + typename IteratorC::Fragment& accum, + int thread_id, + int warp_id, + int lane_id, + cutlass::MatrixCoord const& tile_coords) { + // Non-optimized way to apply exp to registers + // NOTE: accum is attn.T + // TODO: Optimize for each architecture + static constexpr int WarpSize = 32; + using AccumLambdaIterator = + typename DefaultMmaAccumLambdaIterator:: + Iterator; + auto lane_offset = + AccumLambdaIterator::get_lane_offset(lane_id, warp_id, tile_coords); + + AccumLambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) {}, + [&](int accum_m, int accum_n, int idx) { + accum[idx] = expf(accum[idx]); + }, + [&](int accum_m) {}); + accumToSmem(shared_storage, accum, lane_id, tile_coords); + } +}; + +// Simt Specialization +// for f32 on Sm70-Sm75 and f16/f32 below + +template < + int NADim, + typename Operator, + typename OperatorPolicy, + typename scalar_t, + typename WarpShape_, + typename ThreadblockShape_> +struct B2bGemm< + NADim, + cutlass::gemm::warp::MmaSimtTileIterator< + cutlass::MatrixShape<32, 32>, + cutlass::gemm::Operand::kC, + float, + cutlass::layout::RowMajor, + OperatorPolicy, + 1, + 1>, + Operator, + scalar_t, + WarpShape_, + ThreadblockShape_> { + static_assert(NADim >= 1 && NADim < 4); + using Dim = typename natten::cuda::fna::GetDim::type; + + using IteratorC = cutlass::gemm::warp::MmaSimtTileIterator< + cutlass::MatrixShape<32, 32>, + cutlass::gemm::Operand::kC, + float, + cutlass::layout::RowMajor, + OperatorPolicy, + 1, + 1>; + using accum_t = typename IteratorC::Element; + using WarpShape = WarpShape_; + using ThreadblockShape = ThreadblockShape_; + using FragmentC = typename IteratorC::Fragment; + using lse_scalar_t = float; + + // Storage in shared-memory for Q.Kt + using AccumulatorSharedStorage = + cutlass::gemm::threadblock::AccumulatorSharedStorage< + ThreadblockShape, + scalar_t, + cutlass::layout::ColumnMajor, + cutlass::MatrixShape<0, 0> // Padding + >; + + static void CUTLASS_DEVICE accumToSmem( + AccumulatorSharedStorage& shared_storage, + FragmentC const& accum, + int lane_id, + cutlass::MatrixCoord const& tile_coords) { + using Policy = typename IteratorC::Policy; + using Element = typename IteratorC::Element; + using Iterations = typename IteratorC::Iterations; + using Delta = typename IteratorC::Delta; + + auto ref_ = shared_storage.accum_ref(); + // ctor - MmaSimtTileIterator + // compute offset based on thread ID and lane layout + typename Policy::LaneLayout lane_layout = Policy::get_lane_layout(); + + MatrixCoord lane_offset = lane_layout.inverse(lane_id) * + MatrixCoord(Policy::LaneMmaShape::kM, Policy::LaneMmaShape::kN); + + ref_.add_coord_offset(lane_offset); + + // Tile offset + ref_.add_coord_offset( + tile_coords * + cutlass::MatrixCoord( + {IteratorC::Shape::kRow, IteratorC::Shape::kColumn})); + + // store - MmaSimtTileIterator + CUTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Iterations::kColumn; ++mma_n) { + CUTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::LaneMmaShape::kN; ++n) { + CUTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Iterations::kRow; ++mma_m) { + CUTLASS_PRAGMA_UNROLL + for (int m = 0; m < Policy::LaneMmaShape::kM; ++m) { + int r = + Policy::LaneMmaShape::kM * (mma_m * Policy::WarpShape::kRow) + + m; + int c = mma_n * Delta::kColumn + n; + int idx = n + + Policy::LaneMmaShape::kN * + (mma_n + + Iterations::kColumn * + (m + mma_m * Policy::LaneMmaShape::kM)); + ref_.at({r, c}) = scalar_t(accum[idx]); + } + } + } + } + } + + // NOTE(alih): we only apply the exp operator here; + // LSE is pre-fetched into shared memory, and + // subtracted when we check the NA mask and set + // invalid weights to -info. + // Refer to kernel_backward.h for more details. + static void CUTLASS_DEVICE accumApplyExpToSmem( + AccumulatorSharedStorage& shared_storage, + typename IteratorC::Fragment& accum, + int thread_id, + int warp_id, + int lane_id, + cutlass::MatrixCoord const& tile_coords) { + // Non-optimized way to apply exp to registers + // NOTE: accum is attn.T + // TODO: Optimize for each architecture + static constexpr int WarpSize = 32; + using AccumLambdaIterator = + typename DefaultMmaAccumLambdaIterator:: + Iterator; + auto lane_offset = + AccumLambdaIterator::get_lane_offset(lane_id, warp_id, tile_coords); + + AccumLambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) {}, + [&](int accum_m, int accum_n, int idx) { + accum[idx] = expf(accum[idx]); + }, + [&](int accum_m) {}); + accumToSmem(shared_storage, accum, lane_id, tile_coords); + } +}; + +} // namespace threadblock +} // namespace gemm +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fna/gemm/replace_mma_iterators.h b/natten/csrc/include/natten/cuda/fna/gemm/replace_mma_iterators.h new file mode 100644 index 00000000..dbd8f188 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna/gemm/replace_mma_iterators.h @@ -0,0 +1,243 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ + +#pragma once + +#include +#include + +#include +#include + +namespace natten { +namespace cuda { +namespace fna { + +// Technically we can easily replace iterators with a derived +// class, but `MakeCustomMma` dispatches based on Mma classes, +// which makes it really difficult to add arbitrary derived classes, +// so to avoid that, we just replace the Mma iterators by creating +// new Mma instantiations instead of using inheritance. +template +struct ReplaceMmaIterators; + +template < + typename Shape, + typename IteratorA, + typename SmemIteratorA, + cutlass::arch::CacheOperation::Kind CacheOpA, + typename IteratorB, + typename SmemIteratorB, + cutlass::arch::CacheOperation::Kind CacheOpB, + typename ElementC, + typename LayoutC, + typename Policy, + int kStages, + cutlass::gemm::SharedMemoryClearOption SharedMemoryClear, + typename NewIteratorA, + typename NewIteratorB> +struct ReplaceMmaIterators< + cutlass::gemm::threadblock::MmaMultistage< + Shape, + IteratorA, + SmemIteratorA, + CacheOpA, + IteratorB, + SmemIteratorB, + CacheOpB, + ElementC, + LayoutC, + Policy, + kStages, + SharedMemoryClear>, + NewIteratorA, + NewIteratorB> { + using Mma = cutlass::gemm::threadblock::MmaMultistage< + Shape, + NewIteratorA, + SmemIteratorA, + CacheOpA, + NewIteratorB, + SmemIteratorB, + CacheOpB, + ElementC, + LayoutC, + Policy, + kStages, + SharedMemoryClear>; +}; + +template < + typename Shape, + typename IteratorA, + typename SmemIteratorA, + typename IteratorB, + typename SmemIteratorB, + typename ElementC, + typename LayoutC, + typename Policy, + typename NewIteratorA, + typename NewIteratorB> +struct ReplaceMmaIterators< + cutlass::gemm::threadblock::MmaPipelined< + Shape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + Policy>, + NewIteratorA, + NewIteratorB> { + using Mma = cutlass::gemm::threadblock::MmaPipelined< + Shape, + NewIteratorA, + SmemIteratorA, + NewIteratorB, + SmemIteratorB, + ElementC, + LayoutC, + Policy>; +}; + +template +struct ConvertIterators; + +template < + int NADim, + typename Shape, + typename IteratorA, + typename SmemIteratorA, + cutlass::arch::CacheOperation::Kind CacheOpA, + typename IteratorB, + typename SmemIteratorB, + cutlass::arch::CacheOperation::Kind CacheOpB, + typename ElementC, + typename LayoutC, + typename Policy, + int kStages, + cutlass::gemm::SharedMemoryClearOption SharedMemoryClear> +struct ConvertIterators< + NADim, + cutlass::gemm::threadblock::MmaMultistage< + Shape, + IteratorA, + SmemIteratorA, + CacheOpA, + IteratorB, + SmemIteratorB, + CacheOpB, + ElementC, + LayoutC, + Policy, + kStages, + SharedMemoryClear>> { + using NewIteratorA = + cutlass::transform::threadblock::CustomPredicatedTileAccessIterator< + NADim, + typename IteratorA::Shape, + typename IteratorA::Element, + typename IteratorA::Layout, + IteratorA::kAdvanceRank, + typename IteratorA::ThreadMap, + typename IteratorA::AccessType>; + + using NewIteratorB = + cutlass::transform::threadblock::CustomPredicatedTileAccessIterator< + NADim, + typename IteratorB::Shape, + typename IteratorB::Element, + typename IteratorB::Layout, + IteratorB::kAdvanceRank, + typename IteratorB::ThreadMap, + typename IteratorB::AccessType>; + using Mma = cutlass::gemm::threadblock::MmaMultistage< + Shape, + NewIteratorA, + SmemIteratorA, + CacheOpA, + NewIteratorB, + SmemIteratorB, + CacheOpB, + ElementC, + LayoutC, + Policy, + kStages, + SharedMemoryClear>; +}; + +template < + int NADim, + typename Shape, + typename IteratorA, + typename SmemIteratorA, + typename IteratorB, + typename SmemIteratorB, + typename ElementC, + typename LayoutC, + typename Policy> +struct ConvertIterators< + NADim, + cutlass::gemm::threadblock::MmaPipelined< + Shape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + Policy>> { + using NewIteratorA = + cutlass::transform::threadblock::CustomPredicatedTileIterator< + NADim, + typename IteratorA::Shape, + typename IteratorA::Element, + typename IteratorA::Layout, + IteratorA::kAdvanceRank, + typename IteratorA::ThreadMap, + IteratorA::AccessType::kElements>; + + using NewIteratorB = + cutlass::transform::threadblock::CustomPredicatedTileIterator< + NADim, + typename IteratorB::Shape, + typename IteratorB::Element, + typename IteratorB::Layout, + IteratorB::kAdvanceRank, + typename IteratorB::ThreadMap, + IteratorB::AccessType::kElements>; + using Mma = cutlass::gemm::threadblock::MmaPipelined< + Shape, + NewIteratorA, + SmemIteratorA, + NewIteratorB, + SmemIteratorB, + ElementC, + LayoutC, + Policy>; +}; + +} // namespace fna +} // namespace cuda +} // namespace natten diff --git a/natten/csrc/include/natten/cuda/fna/gemm_kernel_utils.h b/natten/csrc/include/natten/cuda/fna/gemm_kernel_utils.h new file mode 100644 index 00000000..5dfc17cd --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna/gemm_kernel_utils.h @@ -0,0 +1,186 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Fused Neighborhood Attention kernels are heavily based on the + * memory-efficient attention kernels from the xFormers project by Meta + * Platforms, Inc. + * + * Copyright (c) Facebook, Inc. and its affiliates + * + * BSD 3-Clause License + * + * 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 names of Facebook, Deepmind Technologies, NYU, NEC + * Laboratories America and IDIAP Research Institute 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 OWNER 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 + +#define DISPATCH_BOOL(BOOL_V, BOOL_NAME, F) \ + { \ + if (BOOL_V) { \ + constexpr bool BOOL_NAME = true; \ + F(); \ + } else { \ + constexpr bool BOOL_NAME = false; \ + F(); \ + } \ + } + +#define CHECK_ALIGNED_PTR(PTR, ALIGNMENT) \ + NATTEN_CHECK(uint64_t(PTR) % ALIGNMENT == 0, #PTR " is not correctly aligned") + +namespace gemm_kernel_utils { + +template +constexpr CUTLASS_HOST_DEVICE integer ceil_div(integer n, integer m) { + return (n + m - 1) / m; +} + +template +constexpr CUTLASS_HOST_DEVICE integer align_up(integer n, integer m) { + return ((n + m - 1) / m) * m; +} + +//////////////////////////////////////////////////////////////////////////////// +// Determine the type of GEMM we do (TensorCores or not, Shapes ...) +// TODO: Maybe we could rely on Cutlass's DefaultGemm templates +//////////////////////////////////////////////////////////////////////////////// + +// Fallback to Simt (FMA on cuda cores) if not in a special case below +template +struct DefaultGemmType { + static constexpr int ThreadK = 8; + static constexpr int WarpK = 8; + static constexpr int kMinimumAlignment = 1; + using InstructionShape = cutlass::gemm::GemmShape<1, 1, 1>; + using OpClass = cutlass::arch::OpClassSimt; + using Operator = cutlass::arch::OpMultiplyAdd; +}; + +// Specialization for tensorcores with f32 +template +struct DefaultGemmType< + ArchTag, + float, + typename cutlass::platform::enable_if< + ArchTag::kMinComputeCapability >= 80>::type> { + static constexpr int ThreadK = 32; + static constexpr int WarpK = 32; + static constexpr int kMinimumAlignment = 4; + using OpClass = cutlass::arch::OpClassTensorOp; + using InstructionShape = cutlass::gemm::GemmShape<16, 8, 8>; + using Operator = cutlass::arch::OpMultiplyAddFastF32; +}; + +// Specialization for tensorcores with f16/bf16 - Sm75+ +template +struct DefaultGemmType< + ArchTag, + scalar_t, + typename cutlass::platform::enable_if< + ArchTag::kMinComputeCapability >= 75 && + cutlass::sizeof_bits::value == 16>::type> { + static constexpr int ThreadK = 32; + static constexpr int WarpK = 32; + static constexpr int kMinimumAlignment = 4; + using OpClass = cutlass::arch::OpClassTensorOp; + using InstructionShape = cutlass::gemm::GemmShape<16, 8, 8>; + using Operator = cutlass::arch::OpMultiplyAdd; +}; + +// Specialization for tensorcores with f16 - Volta +template <> +struct DefaultGemmType { + static constexpr int ThreadK = 32; + static constexpr int WarpK = 32; + static constexpr int kMinimumAlignment = 2; + using OpClass = cutlass::arch::OpClassTensorOp; + using InstructionShape = cutlass::gemm::GemmShape<8, 8, 4>; + using Operator = cutlass::arch::OpMultiplyAdd; +}; + +// Enables to do +// `auto x = kCondition ? fa(arg) : fb(arg)` +// when `fa` and `fb` have different types +template +struct call_conditional; + +template +struct call_conditional { + template + static CUTLASS_HOST_DEVICE auto apply(TA ta, TB tb, Arg arg) + -> decltype(ta(arg)) { + return ta(arg); + } +}; + +template +struct call_conditional { + template + static CUTLASS_HOST_DEVICE auto apply(TA ta, TB tb, Arg arg) + -> decltype(tb(arg)) { + return tb(arg); + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// Mark a variable as warp-uniform - enables some compiler optimizations +// The cheapest way to do it is just to broadcast it from lane 0 +//////////////////////////////////////////////////////////////////////////////// + +template +CUTLASS_DEVICE T warp_uniform(T value) { + struct { + union { + T value; + uint32_t asInt; + }; + } p; + p.value = value; + p.asInt = __shfl_sync(0xffffffff, (unsigned)p.asInt, 0); + return p.value; +} + +template +CUTLASS_DEVICE T* warp_uniform(T* ptr) { + struct { + union { + T* ptr; + uint32_t asInt[2]; + }; + } p; + p.ptr = ptr; + p.asInt[0] = warp_uniform(p.asInt[0]); + p.asInt[1] = warp_uniform(p.asInt[1]); + return p.ptr; +} + +} // namespace gemm_kernel_utils diff --git a/natten/csrc/include/natten/cuda/fna/iterators/default_warp_iterator_from_smem.h b/natten/csrc/include/natten/cuda/fna/iterators/default_warp_iterator_from_smem.h new file mode 100644 index 00000000..a36e1bfb --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna/iterators/default_warp_iterator_from_smem.h @@ -0,0 +1,149 @@ +/* + * Copied from xFormers (https://github.com/facebookresearch/xformers/) + * + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + */ +/*************************************************************************************************** + * Copyright (c) 2017 - 2026 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 Instanciates the right WarpIterator to read from shared memory + The class `DefaultWarpIteratorAFromSharedMemory` is useful when reading + data dumped with `B2bGemm::accumToSmem`. +*/ + +#pragma once + +#include +#include +#include + +#include + +namespace cutlass { +namespace gemm { +namespace threadblock { + +template < + typename WarpShape, + typename InstructionShape, + typename RegularWarpIterator, + typename Policy, + typename Enable = void> +struct DefaultWarpIteratorAFromSharedMemory {}; + +// TensorOp - Ampere half +template +struct DefaultWarpIteratorAFromSharedMemory< + cutlass::gemm::GemmShape<32, 32, 32>, + cutlass::gemm::GemmShape<16, 8, kInstrK>, + RegularWarpIterator, + Policy, + typename platform::enable_if<( + sizeof_bits::value == 16 && + Policy::Operator::Policy::OpDelta::kRow == 1)>::type> { + using OpDelta = typename Policy::Operator::Policy::OpDelta; + using WarpShape = cutlass::MatrixShape<32, 32>; + using InstructionShape = cutlass::gemm::GemmShape<16, 8, kInstrK>; + + using WarpIterator = cutlass::gemm::warp::WarpIteratorFromSmem< + cutlass::gemm::Operand::kA, + typename RegularWarpIterator::Element, + cutlass::MatrixShape>; +}; + +// TensorOp - Ampere f32 +template +struct DefaultWarpIteratorAFromSharedMemory< + WarpShape, + cutlass::gemm::GemmShape<16, 8, 8>, + RegularWarpIterator, + Policy, + typename platform::enable_if<( + sizeof_bits::value != 16 || + Policy::Operator::Policy::OpDelta::kRow != 1)>::type> { + using InstructionShape = cutlass::gemm::GemmShape<16, 8, 8>; + static constexpr auto kWarpSize = 32; + using OpDelta = typename Policy::Operator::Policy::OpDelta; + + using WarpIterator = + cutlass::gemm::warp::MmaTensorOpMultiplicandTileAccessIterator< + cutlass::MatrixShape, + cutlass::gemm::Operand::kA, + typename RegularWarpIterator::Element, + cutlass::layout::RowMajor, + cutlass::MatrixShape, + OpDelta::kRow, + kWarpSize>; +}; + +// TensorOp - Volta +template +struct DefaultWarpIteratorAFromSharedMemory< + WarpShape, + cutlass::gemm::GemmShape<16, 16, 4>, + RegularWarpIterator, + Policy> { + using InstructionShape = cutlass::gemm::GemmShape<16, 16, 4>; + static constexpr auto kWarpSize = 32; + using OpDelta = typename Policy::Operator::Policy::OpDelta; + + using WarpIterator = + cutlass::gemm::warp::MmaVoltaTensorOpMultiplicandTileIterator< + cutlass::MatrixShape<32, 32>, // MatrixShape, + cutlass::gemm::Operand::kA, + typename RegularWarpIterator::Element, + cutlass::layout::RowMajorVoltaTensorOpMultiplicandCrosswise<16, 32>, + cutlass::MatrixShape<16, 4>, + OpDelta::kRow, + kWarpSize>; +}; + +// Simt +template +struct DefaultWarpIteratorAFromSharedMemory< + WarpShape, + cutlass::gemm::GemmShape<1, 1, 1>, + RegularWarpIterator, + Policy> { + using InstructionShape = cutlass::gemm::GemmShape<1, 1, 1>; + static constexpr auto kWarpSize = 32; + + // We just use the same iterator, as we reproduced the same shared-memory + // schema. Just modify it to handle non-complete tiles. + using WarpIterator = RegularWarpIterator; +}; + +} // namespace threadblock +} // namespace gemm +} // namespace cutlass diff --git a/natten/csrc/include/natten/cuda/fna/iterators/epilogue_predicated_tile_iterator.h b/natten/csrc/include/natten/cuda/fna/iterators/epilogue_predicated_tile_iterator.h new file mode 100644 index 00000000..da6b2612 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna/iterators/epilogue_predicated_tile_iterator.h @@ -0,0 +1,562 @@ +/* + * Copied from xFormers (https://github.com/facebookresearch/xformers/) and + * edited. + * + * NOTE: this is only used in the backward pass kernel. + * + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + */ +/*************************************************************************************************** + * Copyright (c) 2017 - 2026 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 Epilogue iterator that supports prefetching + + Mostly copied from "cutlass/epilogue/threadblock/predicated_tile_iterator.h" +*/ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { + +//////////////////////////////////////////////////////////////////////////////// + +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator used to load and store output tile from global memory in +/// epilogue. +/// +/// Satisfies: ReadableTileIterator | PredicatedTileIterator | +/// ForwardTileIterator +/// +template < + int NADim, + typename ThreadMap_, ///< Thread map (conept: OutputTileThreadMap) + typename Element_ ///< Element data type + > +class PredicatedTileIteratorPrefetch { + public: + static_assert(NADim >= 1 && NADim < 4); + using Dim = typename natten::cuda::fna::GetDim::type; + + using ThreadMap = ThreadMap_; + using Shape = typename ThreadMap::Shape; + + using Element = Element_; + + using Layout = layout::RowMajor; + using TensorRef = TensorRef; + using ConstTensorRef = typename TensorRef::ConstTensorRef; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using TensorCoord = MatrixCoord; + + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + static int const kThreads = ThreadMap::kThreads; + static int const kIterations = ThreadMap::Count::kTile; + + static_assert( + ThreadMap::Iterations::kRow > 0, + "ThreadMap::Iterations::kRow must be > 0"); + static_assert( + ThreadMap::Iterations::kGroup > 0, + "ThreadMap::Iterations::kGroup must be > 0"); + static_assert( + ThreadMap::Iterations::kCluster > 0, + "ThreadMap::Iterations::kCluster must be > 0"); + static_assert( + ThreadMap::Iterations::kColumn > 0, + "ThreadMap::Iterations::kColumn must be > 0"); + + /// Fragment object + using Fragment = Array< + Element, + ThreadMap::Iterations::kColumn * ThreadMap::Iterations::kRow * + ThreadMap::Iterations::kGroup * ThreadMap::Iterations::kCluster * + ThreadMap::kElementsPerAccess>; + + /// Memory access size + using AccessType = AlignedArray; + + // + // Parameters struct + // + using ParamsBase = CustomPredicatedTileIteratorParams; + + /// Uses a non-template class + struct Params : ParamsBase { + using Base = ParamsBase; + + CUTLASS_HOST_DEVICE + Params() {} + + CUTLASS_HOST_DEVICE + Params(Dim stride, Dim extent_row) + : ParamsBase( + stride, + int32_t(sizeof(AccessType)) / kElementsPerAccess, + extent_row, + make_OutputTileThreadMapDesc()) {} + + CUTLASS_HOST_DEVICE + Params(Base const& base) : Base(base) {} + }; + + /// Mask object + struct Mask { + static int const kCount = ThreadMap::Iterations::kColumn; + + /// Predicate state + bool predicates[kCount]; + + // + // Mask + // + CUTLASS_HOST_DEVICE + Mask() { + enable(); + } + + ///< Efficiently disables all accesses guarded by mask + CUTLASS_HOST_DEVICE void clear() { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + predicates[i] = false; + } + } + + ///< CUTLASS_HOST_DEVICE enables all accesses guarded by mask + CUTLASS_DEVICE void enable() { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + predicates[i] = true; + } + } + }; + + private: + // + // Data members + // + + /// Parameters structure containing reference and precomputed state. + ParamsBase params_; + + /// Byte-level pointer + uint8_t* byte_pointer_; + + /// Array of boolean values to contain steady-state predicates + Mask mask_; + + /// Extent of the matrix tile in rows + Dim extent_row_; + int32_t extent_row_int; + + /// Extent of the matrix tile in rows + Index extent_column_; + + /// A thread's starting row position (assuming steady-state predicates have + /// been computed) + Index thread_start_row_; + + /// A thread's starting column + Index thread_start_column_; + + /// Internal state counter + int state_[3]; + + // + // Static asserts about internal strides + // + + static_assert(sizeof(extent_column_) == 4, "Expected 32b extents"); + static_assert(sizeof(thread_start_row_) == 4, "Expected 32b extents"); + + private: + // + // Methods + // + + public: + // + // Methods + // + + /// Constructor + CUTLASS_DEVICE + PredicatedTileIteratorPrefetch( + ParamsBase const& params, + Element* pointer, + Dim extent_row, + int32_t extent_col, + int thread_idx, + TensorCoord threadblock_offset = TensorCoord()) + : params_(params) { + TensorCoord thread_offset = + ThreadMap::initial_offset(thread_idx) + threadblock_offset; + + extent_row_ = extent_row; + extent_column_ = extent_col; + extent_row_int = extent_row.prod32(); + + thread_start_row_ = thread_offset.row(); + thread_start_column_ = thread_offset.column(); + + // Initialize predicates + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kColumn; ++c) { + mask_.predicates[c] = + ((thread_offset.column() + ThreadMap::Delta::kColumn * c) < + extent_col); + } + + // Null pointer performs no accesses + if (!pointer) { + mask_.clear(); + } + + // if (ScatterD && !indices) { + // mask_.clear(); + // } + + // Initialize pointer + byte_pointer_ = reinterpret_cast(pointer) + + //(natten::cuda::fna::map_index_to_coord(thread_offset.row(), + // extent_row_) * params_.stride).sum() + + LongIndex(thread_offset.column()) * sizeof(AccessType) / + kElementsPerAccess; + + // if (ScatterD) { + // byte_pointer_ = reinterpret_cast(pointer) + + // LongIndex(thread_offset.column()) * sizeof(AccessType) / + // kElementsPerAccess; + //} + + // Initialize internal state counter + state_[0] = state_[1] = state_[2] = 0; + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + byte_pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_byte_offset(Fragment& frag, int64_t byte_offset) const { + uint8_t* byte_pointer = byte_pointer_; + AccessType* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; + ++cluster) { + CUTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + CUTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + int frag_row_idx = + (row + + ThreadMap::Iterations::kRow * + (group + ThreadMap::Iterations::kGroup * cluster)); + + int32_t row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + + auto row_offset_full_ = natten::cuda::fna::map_index_to_coord( + row_offset + thread_start_row_, extent_row_); + // bool row_guard = + // natten::cuda::fna::is_coord_within_upper_bound(row_offset_full_, + // extent_row_); + bool row_guard = row_offset + thread_start_row_ < extent_row_int; + + AccessType* memory_pointer = reinterpret_cast( + byte_pointer + byte_offset + + (row_offset_full_ * params_.stride).sum()); + + // if (ScatterD && row_guard) { + // assert(indices_); + + // memory_pointer = reinterpret_cast(byte_pointer + + // byte_offset + + // LongIndex(indices_[row_offset + thread_start_row_]) * + // LongIndex(params_.stride)); + //} + + CUTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; + ++column) { + bool guard = row_guard && mask_.predicates[column]; + + cutlass::arch::global_load( + frag_ptr + [frag_row_idx * ThreadMap::Iterations::kColumn + column], + (void*)&memory_pointer + [column * ThreadMap::Delta::kColumn / kElementsPerAccess], + guard); + } + + if (row + 1 < ThreadMap::Iterations::kRow) { + // if (!ScatterD) { + // byte_pointer += params_.increment_row; + //} + } + } + + if (group + 1 < ThreadMap::Iterations::kGroup) { + // byte_pointer += params_.increment_group; + } + } + + if (cluster + 1 < ThreadMap::Iterations::kCluster) { + // byte_pointer += params_.increment_cluster; + } + } + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) const { + load_with_byte_offset(frag, 0); + } + + /// Stores a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, int64_t byte_offset) const { + uint8_t* byte_pointer = byte_pointer_; + AccessType const* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; + ++cluster) { + CUTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + CUTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + int frag_row_idx = + (row + + ThreadMap::Iterations::kRow * + (group + ThreadMap::Iterations::kGroup * cluster)); + + int32_t row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + + auto row_offset_full_ = natten::cuda::fna::map_index_to_coord( + row_offset + thread_start_row_, extent_row_); + // bool row_guard = + // natten::cuda::fna::is_coord_within_upper_bound(row_offset_full_, + // extent_row_); + bool row_guard = row_offset + thread_start_row_ < extent_row_int; + + AccessType* memory_pointer = reinterpret_cast( + byte_pointer + byte_offset + + (row_offset_full_ * params_.stride).sum()); + + // if (ScatterD && row_guard) { + // assert(indices_); + + // memory_pointer = reinterpret_cast(byte_pointer + + // byte_offset + + // LongIndex(indices_[row_offset + thread_start_row_]) * + // LongIndex(params_.stride)); + //} + + CUTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; + ++column) { + bool guard = row_guard && mask_.predicates[column]; + + // if (UseCUDAStore) { + // if (guard) { + // memory_pointer + // [column * ThreadMap::Delta::kColumn / kElementsPerAccess] + // = + // frag_ptr + // [frag_row_idx * ThreadMap::Iterations::kColumn + + // column]; + // } + // } else { + cutlass::arch::global_store( + frag_ptr + [frag_row_idx * ThreadMap::Iterations::kColumn + column], + (void*)&memory_pointer[0], + //(void*)&memory_pointer + // [column * ThreadMap::Delta::kColumn / kElementsPerAccess], + guard); + //} + memory_pointer += (ThreadMap::Delta::kColumn / kElementsPerAccess); + } + + if (row + 1 < ThreadMap::Iterations::kRow) { + // if (!ScatterD) { + // byte_pointer += params_.increment_row; + //} + } + } + + if (group + 1 < ThreadMap::Iterations::kGroup) { + // byte_pointer += params_.increment_group; + } + } + + if (cluster + 1 < ThreadMap::Iterations::kCluster) { + // byte_pointer += params_.increment_cluster; + } + } + } + + /// Stores a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) const { + store_with_byte_offset(frag, 0); + } + + CUTLASS_DEVICE + MatrixCoord thread_start() const { + return MatrixCoord(thread_start_row_, thread_start_column_); + } + + /// Need to get the thread start row from the tile iterator + CUTLASS_DEVICE + int32_t thread_start_row() const { + return thread_start_row_; + } + + /// Need to get the thread start row from the tile iterator + CUTLASS_DEVICE + int32_t thread_start_column() const { + return thread_start_column_; + } + + /// Extent of the matrix in columns + CUTLASS_DEVICE + Index extent_column() const { + return extent_column_; + } + + /// Advances to the next position to load or store + CUTLASS_HOST_DEVICE + PredicatedTileIteratorPrefetch& operator++() { + ++state_[0]; + + // if (!ScatterD) { + // byte_pointer_ += params_.advance_row; + //} + + thread_start_row_ += ThreadMap::Shape::kRow; + + if (state_[0] == ThreadMap::Count::kRow) { + state_[0] = 0; + ++state_[1]; + // byte_pointer_ += params_.advance_group; + + thread_start_row_ += (ThreadMap::Shape::kGroup - 1) * + ThreadMap::Shape::kRow * ThreadMap::Count::kRow; + + if (state_[1] == ThreadMap::Count::kGroup) { + state_[1] = 0; + ++state_[2]; + // byte_pointer_ += params_.advance_cluster; + + thread_start_row_ += ThreadMap::Count::kGroup * + ThreadMap::Shape::kGroup * ThreadMap::Count::kRow * + ThreadMap::Shape::kRow; + + if (state_[2] == ThreadMap::Count::kCluster) { + state_[2] = 0; + // byte_pointer_ += params_.advance_tile; + } + } + } + + return *this; + } + + ///< Efficiently disables all accesses guarded by mask + CUTLASS_DEVICE void clear_mask() { + mask_.clear(); + } + + ///< Efficiently enables all accesses guarded by mask + CUTLASS_DEVICE void enable_mask() { + mask_.enable(); + } + + ///< Sets the mask + CUTLASS_DEVICE void get_mask(Mask& mask) const { + mask = mask_; + } + + ///< Sets the mask + CUTLASS_DEVICE void set_mask(Mask const& mask) { + mask_ = mask; + } +}; + +template +struct MakePrefetchableIterator { + using Iterator = PredicatedTileIteratorPrefetch< + NADim, + typename IT::ThreadMap, + typename IT::Element>; +}; + +/////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fna/iterators/make_residual_last.h b/natten/csrc/include/natten/cuda/fna/iterators/make_residual_last.h new file mode 100644 index 00000000..14a404aa --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna/iterators/make_residual_last.h @@ -0,0 +1,108 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Fused Neighborhood Attention kernels are heavily based on the + * memory-efficient attention kernels from the xFormers project by Meta + * Platforms, Inc. + * + * Copyright (c) Facebook, Inc. and its affiliates + * + * BSD 3-Clause License + * + * 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 names of Facebook, Deepmind Technologies, NYU, NEC + * Laboratories America and IDIAP Research Institute 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 OWNER 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 + +namespace cutlass { +namespace transform { +namespace threadblock { + +template +struct MakeIteratorResidualLast; + +template < + int NADim, + typename Shape, + typename Element, + typename Layout, + int AdvanceRank, + typename ThreadMap, + int AccessSize> +struct MakeIteratorResidualLast< + NADim, + PredicatedTileIterator< + Shape, + Element, + Layout, + AdvanceRank, + ThreadMap, + AccessSize>> { + using Iterator = PredicatedTileIteratorResidualLast< + NADim, + Shape, + Element, + Layout, + AdvanceRank, + ThreadMap, + AccessSize>; +}; + +template < + int NADim, + typename Shape, + typename Element, + typename Layout, + int AdvanceRank, + typename ThreadMap, + typename AccessType> +struct MakeIteratorResidualLast< + NADim, + PredicatedTileAccessIterator< + Shape, + Element, + Layout, + AdvanceRank, + ThreadMap, + AccessType>> { + using Iterator = PredicatedTileAccessIteratorResidualLast< + NADim, + Shape, + Element, + Layout, + AdvanceRank, + ThreadMap, + AccessType>; +}; + +} // namespace threadblock +} // namespace transform +} // namespace cutlass diff --git a/natten/csrc/include/natten/cuda/fna/iterators/predicated_tile_access_iterator.h b/natten/csrc/include/natten/cuda/fna/iterators/predicated_tile_access_iterator.h new file mode 100644 index 00000000..c76032b2 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna/iterators/predicated_tile_access_iterator.h @@ -0,0 +1,1662 @@ +/* + * Copied from CUTLASS (https://github.com/NVIDIA/cutlass/) and edited. + */ +/*************************************************************************************************** + * Copyright (c) 2017 - 2026 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 Templates calculating the address and predicates to the load of tiles + from pitch-linear rank=2 tensors. + + This iterator uses masks to guard out-of-bounds accesses. The first tile + this iterator visits maybe partial, then the remaining tiles are complete. + So, we only need to compute the predicates twice, once before the first tile + and once for the remaining full tiles which can share the same predicates. + + A precomputed "Params" object minimizes the amount of state that must be + stored in registers, and integer addition is used to advance the pointer + through memory. +*/ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +//////////////////////////////////////////////////////////////////////////////// + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace transform { +namespace threadblock { + +template +struct CustomPredicatedTileAccessIteratorParams { + using Dim = typename natten::cuda::fna::GetDim::type; + + using Index = int32_t; + using LongIndex = + typename std::conditional::type; + + // + // Data members + // + /// stride of pitch-linear layout (units of Element) + Dim stride_; + /// amount (in byte) to increment pointer to move to next access along + /// strided dimension + LongIndex inc_strided_; + /// amount (in byte) to increment pointer from last access to first access + /// of next tile + LongIndex inc_next_; + /// amount (in byte) to increment pointer from first access of current tile + /// to first access of next tile + LongIndex inc_advance_; + + // + // Methods + // + + CUTLASS_HOST_DEVICE + Status initialize( + Dim stride, + Dim extent_row, + PredicatedTileAccessIteratorDesc desc) { + stride_ = stride; + + if constexpr (NADim == 1) { + inc_strided_ = (LongIndex(stride_.x) * desc.threadmap_delta.strided()) * + desc.element_size_bits / 8; + + if (desc.advance_rank) { + // advance along strided dimension + inc_advance_ = desc.threadblock_shape.strided() * LongIndex(stride_.x) * + desc.element_size_bits / 8; + } else { + // advance along contiguous dimension + inc_advance_ = + desc.threadblock_shape.contiguous() * desc.element_size_bits / 8; + } + + inc_next_ = inc_advance_ - + LongIndex(desc.threadmap_iterations.strided() - 1) * + desc.threadmap_delta.strided() * LongIndex(stride_.x) * + desc.element_size_bits / 8; + // inc_strided_ = + // (natten::cuda::fna::map_index_to_coord(desc.threadmap_delta.strided(), + // extent_row) * stride_).sum() * + // desc.element_size_bits / 8; + //// inc_strided_ = (LongIndex(stride_) * desc.threadmap_delta.strided()) + ///* / desc.element_size_bits / 8; + + // if (desc.advance_rank) { + // // advance along strided dimension + // inc_advance_ = (natten::cuda::fna::map_index_to_coord( + // desc.threadblock_shape.strided(), extent_row) * stride_).sum() * + // desc.element_size_bits / 8; + // // inc_advance_ = + // // desc.threadblock_shape.strided() * LongIndex(stride_) * + // desc.element_size_bits / 8; + //} else { + // // advance along contiguous dimension + // inc_advance_ = desc.threadblock_shape.contiguous() * + // desc.element_size_bits / 8; + //} + + // inc_next_ = inc_advance_ - + // (natten::cuda::fna::map_index_to_coord((desc.threadmap_iterations.strided() + // - 1) * + // desc.threadmap_delta.strided(), + // extent_row) * stride_).sum() * + // desc.element_size_bits / 8; + //// inc_next_ = inc_advance_ - + /// LongIndex(desc.threadmap_iterations.strided() - 1) * / + /// desc.threadmap_delta.strided() * LongIndex(stride_) * / + /// desc.element_size_bits / 8; + } else { + // inc_strided_ = LongIndex(desc.threadmap_delta.strided()); + + // if (desc.advance_rank) { + // // advance along strided dimension + // inc_advance_ = LongIndex(desc.threadblock_shape.strided()); + // } else { + // // advance along contiguous dimension + // inc_advance_ = LongIndex(desc.threadblock_shape.contiguous()); + // } + + // inc_next_ = inc_advance_ - + // LongIndex(desc.threadmap_iterations.strided() - 1) * + // desc.threadmap_delta.strided(); + } + + return Status::kSuccess; + } + + CUTLASS_HOST_DEVICE + CustomPredicatedTileAccessIteratorParams( + Dim stride, + Dim extent_row, + PredicatedTileAccessIteratorDesc desc) { + initialize(stride, extent_row, desc); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// CustomPredicatedTileAccessIteratorPredicates +/// +template < + int NADim, + typename Shape_, + typename Element_, + typename Layout_, + int AdvanceRank, + typename ThreadMap_, + typename AccessType_> +class CustomPredicatedTileAccessIteratorPredicates { + public: + static_assert(NADim >= 1 && NADim < 4); + using Dim = typename natten::cuda::fna::GetDim::type; + + using Shape = Shape_; + using Element = Element_; + using Layout = Layout_; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorCoord = typename Layout::TensorCoord; + + static int const kAccessesPerVector = + ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert( + !(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + static int const kPredicatesPerByte = 4; + static int const kPredicatesPerWord = 4 * kPredicatesPerByte; + + static int const kPredicateCount = + ThreadMap::Iterations::kCount * kAccessesPerVector; + + /// Number of 32b words containing predicates + static int const kPredicateByteCount = + (kPredicateCount + kPredicatesPerByte - 1) / kPredicatesPerByte; + static int const kPredicateWordCount = (kPredicateByteCount + 3) / 4; + + static unsigned const kPredicateMask = (1u << kPredicatesPerByte) - 1u; + + static_assert(kPredicateWordCount <= 4, "Too many predicates."); + + /// Predicate vector stores mask to guard accesses + using Mask = Array; + + // private: + /// Guard predicates + uint32_t predicates_[kPredicateWordCount]; + + /// Size of tensor + Dim extent_row_dim; + TensorCoord extent_; + + /// Initial offset for each thread + TensorCoord thread_offset_; + + /// Offset to the first steady-state tile + TensorCoord residue_offset_; + + /// Iteration along vectors implied by the thread map + int iteration_vector_; + + /// Iteration in the contiguous dimension + int iteration_contiguous_; + + /// Iteration in the strided dimension + int iteration_strided_; + + public: + /// Computes predicates based on internally tracked per-thread offset. + CUTLASS_DEVICE + void compute_predicates_( + /// Extent of the matrix window + TensorCoord extent, + /// optionally, simplify predicate calculation during 'steady state' phase + bool is_steady_state = false) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPredicateWordCount; ++i) { + predicates_[i] = 0u; + } + + CUTLASS_PRAGMA_UNROLL + for (int access_idx = 0; + access_idx < ThreadMap::Iterations::kCount * kAccessesPerVector; + ++access_idx) { + int s = access_idx / + (ThreadMap::Iterations::kContiguous * kAccessesPerVector); + + int access_residual = access_idx % + (ThreadMap::Iterations::kContiguous * kAccessesPerVector); + + int c = access_residual / kAccessesPerVector; + int v = access_residual % kAccessesPerVector; + + TensorCoord iteration_coord( + c * ThreadMap::Delta::kContiguous + v * AccessType::kElements, + s * ThreadMap::Delta::kStrided); + + TensorCoord coord = thread_offset_ + iteration_coord; + + bool guard; + + // Dim coord_row; + // if constexpr (kAdvanceRank == 0) { + // coord_row = natten::cuda::fna::map_index_to_coord(coord.strided(), + // extent_.strided()); + //} else { + // coord_row = natten::cuda::fna::map_index_to_coord(coord.contiguous(), + // natten::cuda::fna::MapTileSizeToNd::value); + //} + + // if (is_steady_state) { + // guard = coord.strided() < extent.strided(); + // } else { + ////if constexpr (kAdvanceRank == 0) { + // guard = (coord.strided() < extent.strided() && + // coord.contiguous() < extent.contiguous()); + ////} else { + //// guard = (coord.strided() < extent_col && + //// coord_row < extent_row); + ////} + //} + if (is_steady_state) { + if (kAdvanceRank == 0) { + guard = (coord.strided() < extent.strided()); + } else { + guard = (coord.contiguous() < extent.contiguous()); + } + } else { + guard = + (coord.strided() < extent.strided() && + coord.contiguous() < extent.contiguous()); + } + + int pred_idx = + v + kAccessesPerVector * (c + ThreadMap::Iterations::kContiguous * s); + + int word_idx = pred_idx / kPredicatesPerWord; + int residual = pred_idx % kPredicatesPerWord; + int byte_idx = residual / kPredicatesPerByte; + int bit_idx = residual % kPredicatesPerByte; + + predicates_[word_idx] |= (unsigned(guard) << (byte_idx * 8 + bit_idx)); + } + } + + CUTLASS_HOST_DEVICE + void set_predicates(int thread_id, TensorCoord const& threadblock_offset) { + TensorCoord residue_extent; + // int32_t residue_extent_row = 0; + // int32_t residue_extent_col = 0; + // if constexpr (kAdvanceRank == 0) { + // // Column-major right handside operand (only instance is K in QK^T). + // // extent_[0] == K (dim per head) == extent_.contiguous() + // // extent_.contiguous() == extent_[0] == extent_.contiguous() + // // + // // Row-major left handside operand (only instance is Q in QK^T) + // // extent_[0] == K (dim per head) == extent_.contiguous() + // // extent_.contiguous() == extent_[0] == extent_.contiguous() + // typename TensorCoord::Index residue_size = (extent_.contiguous() - + // threadblock_offset.contiguous()) % Shape::kContiguous; if + // (!residue_size) { + // residue_size = Shape::kContiguous; + // } + // residue_offset_ = make_Coord(residue_size, 0); + // + // residue_extent.strided() = extent_.strided(); + // residue_extent.contiguous() = min(extent_.contiguous(), + // threadblock_offset.contiguous() + residue_size); + //} else { + // // Row-major right-handside operand (only instance is V in PV) + // // extent_[1] == M == extent_.strided() + // // extent_.contiguous() == extent_[0] == extent_.contiguous() + // // extent_.strided() == extent_[1] == extent_.strided() + // // + // typename TensorCoord::Index residue_size = (extent_.strided() - + // threadblock_offset.strided()) % Shape::kStrided; if (!residue_size) { + // residue_size = Shape::kStrided; + // } + // residue_offset_ = make_Coord(0, residue_size); + + // residue_extent.strided() = min(threadblock_offset.strided() + + // residue_size, extent_.strided()); residue_extent.contiguous() = + // extent_.contiguous(); + //} + if (kAdvanceRank) { + typename TensorCoord::Index residue_size = + (extent_[kAdvanceRank] - threadblock_offset.strided()) % + Shape::kStrided; + if (!residue_size) { + residue_size = Shape::kStrided; + } + + residue_offset_ = make_Coord(0, residue_size); + residue_extent = make_Coord( + extent_.contiguous(), + min(threadblock_offset.strided() + residue_size, extent_.strided())); + } else { + typename TensorCoord::Index residue_size = + (extent_[kAdvanceRank] - threadblock_offset.contiguous()) % + Shape::kContiguous; + if (!residue_size) { + residue_size = Shape::kContiguous; + } + + residue_offset_ = make_Coord(residue_size, 0); + + residue_extent = make_Coord( + min(extent_.contiguous(), + threadblock_offset.contiguous() + residue_size), + extent_.strided()); + } + + // Per-thread offset in logical coordinates of tensor + thread_offset_ = threadblock_offset + ThreadMap::initial_offset(thread_id); + + compute_predicates_(residue_extent, false); + + set_iteration_index(0); + } + + /// Default constructor + CustomPredicatedTileAccessIteratorPredicates() = default; + + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + CustomPredicatedTileAccessIteratorPredicates( + /// Extent of tensor (rows) + Dim extent_row, + /// Extent of tensor (columns) + int32_t extent_col) + : extent_row_dim(extent_row), extent_(extent_col, extent_row.prod32()) {} + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) { + iteration_vector_ = index % kAccessesPerVector; + int residual_access = index / kAccessesPerVector; + + iteration_contiguous_ = + residual_access % ThreadMap::Iterations::kContiguous; + iteration_strided_ = residual_access / ThreadMap::Iterations::kContiguous; + } + + /// Increment and return an instance to self. + CUTLASS_HOST_DEVICE + CustomPredicatedTileAccessIteratorPredicates& operator++() { + return *this; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPredicateWordCount; ++i) { + predicates_[i] = enable ? 0u : predicates_[i]; + } + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPredicateWordCount; ++i) { + predicates_[i] = 0xffffffff; + } + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPredicateWordCount; ++i) { + predicates_[i] = mask[i]; + } + } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPredicateWordCount; ++i) { + mask[i] = predicates_[i]; + } + } + + /// Returns whether access is valid or not + CUTLASS_HOST_DEVICE + bool valid() const { + int pred_idx = iteration_vector_ + + kAccessesPerVector * + (iteration_contiguous_ + + iteration_strided_ * ThreadMap::Iterations::kContiguous); + + int word_idx = pred_idx / kPredicatesPerWord; + int residual = pred_idx % kPredicatesPerWord; + int byte_idx = residual / kPredicatesPerByte; + int bit_idx = residual % kPredicatesPerByte; + + bool pred = (predicates_[word_idx] & (1u << (byte_idx * 8 + bit_idx))) != 0; + return pred; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// CustomPredicatedTileAccessIterator +/// +template < + int NADim, + typename Shape, + typename Element, + typename Layout, + int AdvanceRank, + typename ThreadMap, + typename AccessType> +class CustomPredicatedTileAccessIterator; + +//////////////////////////////////////////////////////////////////////////////// + +/// CustomPredicatedTileAccessIterator for possibly multiple axes (NA1D, 2D, 3D) +template < + int NADim, + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + typename AccessType_> +class CustomPredicatedTileAccessIterator< + NADim, + Shape_, + Element_, + layout::PitchLinear, + AdvanceRank, + ThreadMap_, + AccessType_> { + public: + static_assert(NADim >= 1 && NADim < 4); + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Dim = typename natten::cuda::fna::GetDim::type; + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::PitchLinear; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingPredicates = CustomPredicatedTileAccessIteratorPredicates< + NADim, + Shape, + Element, + Layout, + AdvanceRank, + ThreadMap, + AccessType>; + + static int const kAccessesPerVector = + ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert( + !(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + using Mask = typename UnderlyingPredicates::Mask; + + /// Uses a non-template class + struct Params : CustomPredicatedTileAccessIteratorParams { + using Base = CustomPredicatedTileAccessIteratorParams; + + /// Default constructor + Params() = default; + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Dim stride, Dim extent_row) + : Base( + stride, + extent_row, + MakePredicatedTileAccessIteratorDesc< + Shape, + Element, + Layout, + kAdvanceRank, + ThreadMap>()()) {} + + CUTLASS_HOST_DEVICE + Params(Base const& base) : Base(base) {} + }; + + private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char*; + + private: + // + // Data members + // + + UnderlyingPredicates the_predicates; + + /// Parameters object with precomputed internal state + Params params_; + + /// Internal pointer to first access of tile + BytePointer pointer_; + TensorCoord coord_offset_; + + /// Used for out-of-order visitation + bool is_residue_tile_; + + private: + /// Computes predicates based on internally tracked per-thread offset. + CUTLASS_DEVICE + void compute_predicates_( + /// Extent of the matrix window + TensorCoord extent, + /// optionally, simplify predicate calculation during 'steady state' phase + bool is_steady_state = false) { + the_predicates.compute_predicates_(extent, is_steady_state); + } + + public: + /// Default constructor + CustomPredicatedTileAccessIterator() = default; + + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + CustomPredicatedTileAccessIterator( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor (rows) + Dim extent_row, + /// Extent of tensor (columns) + int32_t extent_col, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const& threadblock_offset) + : params_(params), + pointer_(reinterpret_cast( + const_cast(pointer))), + the_predicates(extent_row, extent_col), + is_residue_tile_(true) { + the_predicates.set_predicates(thread_id, threadblock_offset); + + // update internal pointers + coord_offset_ = the_predicates.thread_offset_; + // add_pointer_offset(the_predicates.thread_offset_.strided(), + // the_predicates.thread_offset_.contiguous()); + } + + /// Construct a CustomPredicatedTileAccessIterator with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + CustomPredicatedTileAccessIterator( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor (rows) + Dim extent_row, + /// Extent of tensor (columns) + int32_t extent_col, + ///< ID of each participating thread + int thread_id) + : CustomPredicatedTileAccessIterator( + params, + pointer, + extent_row, + extent_col, + thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) { + the_predicates.set_iteration_index(index); + } + + CUTLASS_HOST_DEVICE + void add_pointer_offset(int32_t row_offset, int32_t col_offset) { + coord_offset_.strided() += row_offset; + coord_offset_.contiguous() += col_offset; + } + + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex linear_offset) { + printf("FATAL!"); + // if constexpr (kAdvanceRank) { + // auto col_offset = linear_offset % Shape::kStrided; + // auto row_offset = linear_offset / Shape::kStrided; + // add_pointer_offset(row_offset, (int32_t)col_offset); + // } else { + // auto col_offset = linear_offset % Shape::kContiguous; + // auto row_offset = linear_offset / Shape::kContiguous; + // add_pointer_offset(row_offset, (int32_t)col_offset); + // } + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_DEVICE + void add_tile_offset(TensorCoord const& tile_offset) { + if (is_residue_tile_) { + the_predicates.thread_offset_ += the_predicates.residue_offset_; + + the_predicates.compute_predicates_(the_predicates.extent_, true); + + // add_pointer_offset(the_predicates.residue_offset_.strided(), + // the_predicates.residue_offset_.contiguous()); + + coord_offset_.strided() = the_predicates.thread_offset_.strided() + + Shape::kStrided * (tile_offset.strided() - kAdvanceRank); + coord_offset_.contiguous() = the_predicates.thread_offset_.contiguous() + + Shape::kContiguous * (tile_offset.contiguous() - (1 - kAdvanceRank)); + + // if (kAdvanceRank) { + // add_pointer_offset( + // params_.inc_advance_ * (tile_offset.strided() - 1) + + // Shape::kContiguous * tile_offset.contiguous()); + // // pointer_ += params_.inc_advance_ * LongIndex(tile_offset.strided() + // - 1); + // // pointer_ += Shape::kContiguous * tile_offset.contiguous() * + // sizeof_bits::value / 8; + // } else { + // // Advance along contiguous dimension + // add_pointer_offset( + // params_.inc_advance_ * (tile_offset.contiguous() - 1) + + // Shape::kStrided * tile_offset.strided()); + // // pointer_ += params_.inc_advance_ * + // LongIndex(tile_offset.contiguous() - 1); + // // pointer_ += Shape::kStrided * tile_offset.strided() * + // sizeof_bits::value / 8; + // } + } else { + add_pointer_offset( + Shape::kStrided * tile_offset.strided(), + Shape::kContiguous * tile_offset.contiguous()); + // if (kAdvanceRank) { + // add_pointer_offset( + // params_.inc_advance_ * tile_offset.strided() + + // Shape::kContiguous * tile_offset.contiguous()); + // // pointer_ += params_.inc_advance_ * + // LongIndex(tile_offset.strided()); + // // pointer_ += Shape::kContiguous * tile_offset.contiguous(); + // } else { + // // Advance along contiguous dimension + // add_pointer_offset( + // params_.inc_advance_ * tile_offset.contiguous() + + // Shape::kStrided * tile_offset.strided()); + // // pointer_ += params_.inc_advance_ * + // LongIndex(tile_offset.contiguous()); + // // pointer_ += Shape::kStrided * tile_offset.strided(); + // } + } + + is_residue_tile_ = false; + } + + CUTLASS_HOST_DEVICE + LongIndex get_pointer_offset() const { + LongIndex coord_contig = coord_offset_.contiguous() + + the_predicates.iteration_contiguous_ * ThreadMap::Delta::kContiguous; + LongIndex coord_strided = + (natten::cuda::fna::map_index_to_coord( + coord_offset_.strided() + + the_predicates.iteration_strided_ * ThreadMap::Delta::kStrided, + the_predicates.extent_row_dim) * + params_.stride_) + .sum(); + return ((coord_strided + coord_contig) * sizeof_bits::value) / 8; + } + + /// Returns a pointer + CUTLASS_HOST_DEVICE + AccessType* get() const { + if (!valid()) { + return nullptr; + } + return reinterpret_cast(pointer_ + get_pointer_offset()) + + the_predicates.iteration_vector_; + } + + /// Increment and return an instance to self. + CUTLASS_HOST_DEVICE + CustomPredicatedTileAccessIterator& operator++() { + the_predicates.operator++(); + + ++the_predicates.iteration_vector_; + if (the_predicates.iteration_vector_ < kAccessesPerVector) { + return *this; + } + + the_predicates.iteration_vector_ = 0; + ++the_predicates.iteration_contiguous_; + + if (the_predicates.iteration_contiguous_ < + ThreadMap::Iterations::kContiguous) { + return *this; + } + + // Enter here only if (iteration_contiguous_ == + // ThreadMap::Iteration::kContiguous) + the_predicates.iteration_contiguous_ = 0; + ++the_predicates.iteration_strided_; + + if (the_predicates.iteration_strided_ < ThreadMap::Iterations::kStrided) { + // add_pointer_offset(params_.inc_strided_, 0); + // pointer_ += params_.inc_strided_; + + return *this; + } + + // Enter here only if (iteration_stride_ == ThreadMap::Iteration::kStrided) + // which means we enter the next tile. + the_predicates.iteration_strided_ = 0; + + // // advance to next tile + // add_pointer_offset(params_.inc_next_, 0); + // // pointer_ += params_.inc_next_; + + // //if constexpr (kAdvanceRank) { + // add_pointer_offset(-1 * params_.inc_advance_); + // //} else { + // // // along the contiguous axis. + // // add_pointer_offset(0, -1 * params_.inc_advance_); + // //} + // // now return to start tile - if the iterator is subsequently advanced, + // this + // // subtraction as well as the subsequent integer addition are both elided + // by + // // the compiler. + // // pointer_ -= params_.inc_advance_; + + return *this; + } + + /// Increment and return an instance to self. + CUTLASS_HOST_DEVICE + CustomPredicatedTileAccessIterator operator++(int) { + CustomPredicatedTileAccessIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + the_predicates.clear_mask(enable); + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { + the_predicates.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { + the_predicates.set_mask(mask); + } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { + the_predicates.get_mask(mask); + } + + /// Returns whether access is valid or not + CUTLASS_HOST_DEVICE + bool valid() const { + return the_predicates.valid(); + } +}; + +// Specialization for NA1D, where we only have one stride, +// and therefore the typical pitch-linear layout, and won't have +// to manually convert offsets on every get(). +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + typename AccessType_> +class CustomPredicatedTileAccessIterator< + 1, + Shape_, + Element_, + layout::PitchLinear, + AdvanceRank, + ThreadMap_, + AccessType_> { + public: + static constexpr int NADim = 1; + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Dim = typename natten::cuda::fna::GetDim::type; + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::PitchLinear; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingPredicates = CustomPredicatedTileAccessIteratorPredicates< + NADim, + Shape, + Element, + Layout, + AdvanceRank, + ThreadMap, + AccessType>; + + static int const kAccessesPerVector = + ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert( + !(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + using Mask = typename UnderlyingPredicates::Mask; + + /// Uses a non-template class + struct Params : CustomPredicatedTileAccessIteratorParams { + using Base = CustomPredicatedTileAccessIteratorParams; + + /// Default constructor + Params() = default; + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Dim stride, Dim extent_row) + : Base( + stride, + extent_row, + MakePredicatedTileAccessIteratorDesc< + Shape, + Element, + Layout, + kAdvanceRank, + ThreadMap>()()) {} + + CUTLASS_HOST_DEVICE + Params(Base const& base) : Base(base) {} + }; + + private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char*; + + private: + // + // Data members + // + + UnderlyingPredicates the_predicates; + + /// Parameters object with precomputed internal state + Params params_; + + /// Internal pointer to first access of tile + BytePointer pointer_; + + /// Used for out-of-order visitation + bool is_residue_tile_; + + private: + /// Computes predicates based on internally tracked per-thread offset. + CUTLASS_DEVICE + void compute_predicates_( + /// Extent of the matrix window + Dim extent_row, + int32_t extent_col, + /// optionally, simplify predicate calculation during 'steady state' phase + bool is_steady_state = false) { + the_predicates.compute_predicates_(extent_row, extent_col, is_steady_state); + } + + public: + /// Default constructor + CustomPredicatedTileAccessIterator() = default; + + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + CustomPredicatedTileAccessIterator( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor (rows) + Dim extent_row, + /// Extent of tensor (columns) + int32_t extent_col, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const& threadblock_offset) + : params_(params), + pointer_(reinterpret_cast( + const_cast(pointer))), + the_predicates(extent_row, extent_col), + is_residue_tile_(true) { + the_predicates.set_predicates(thread_id, threadblock_offset); + + // update internal pointers + auto offset = (natten::cuda::fna::map_index_to_coord( + the_predicates.thread_offset_.strided(), + the_predicates.extent_row_dim) * + params_.stride_) + .sum() + + the_predicates.thread_offset_.contiguous(); + + add_pointer_offset(offset); + } + + /// Construct a CustomPredicatedTileAccessIterator with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + CustomPredicatedTileAccessIterator( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor (rows) + Dim extent_row, + /// Extent of tensor (columns) + int32_t extent_col, + ///< ID of each participating thread + int thread_id) + : CustomPredicatedTileAccessIterator( + params, + pointer, + extent_row, + extent_col, + thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) { + the_predicates.set_iteration_index(index); + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += sizeof_bits::value * pointer_offset / 8; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_DEVICE + void add_tile_offset(TensorCoord const& tile_offset) { + if (is_residue_tile_) { + the_predicates.thread_offset_ += the_predicates.residue_offset_; + + the_predicates.compute_predicates_(the_predicates.extent_, true); + + auto offset = (natten::cuda::fna::map_index_to_coord( + the_predicates.residue_offset_.strided(), + the_predicates.extent_row_dim) * + params_.stride_) + .sum() + + the_predicates.residue_offset_.contiguous(); + + add_pointer_offset(offset); + + if (kAdvanceRank) { + pointer_ += params_.inc_advance_ * LongIndex(tile_offset.strided() - 1); + pointer_ += Shape::kContiguous * tile_offset.contiguous() * + sizeof_bits::value / 8; + } else { + pointer_ += + params_.inc_advance_ * LongIndex(tile_offset.contiguous() - 1); + pointer_ += Shape::kStrided * tile_offset.strided() * + sizeof_bits::value / 8; + } + } else { + if (kAdvanceRank) { + pointer_ += params_.inc_advance_ * LongIndex(tile_offset.strided()); + pointer_ += Shape::kContiguous * tile_offset.contiguous(); + } else { + pointer_ += params_.inc_advance_ * LongIndex(tile_offset.contiguous()); + pointer_ += Shape::kStrided * tile_offset.strided(); + } + } + + is_residue_tile_ = false; + } + + /// Returns a pointer + CUTLASS_HOST_DEVICE + AccessType* get() const { + return reinterpret_cast( + pointer_ + + the_predicates.iteration_contiguous_ * + (ThreadMap::Delta::kContiguous * + sizeof_bits::value) / + 8) + + the_predicates.iteration_vector_; + } + + /// Increment and return an instance to self. + CUTLASS_HOST_DEVICE + CustomPredicatedTileAccessIterator& operator++() { + the_predicates.operator++(); + + ++the_predicates.iteration_vector_; + if (the_predicates.iteration_vector_ < kAccessesPerVector) { + return *this; + } + + the_predicates.iteration_vector_ = 0; + ++the_predicates.iteration_contiguous_; + + if (the_predicates.iteration_contiguous_ < + ThreadMap::Iterations::kContiguous) { + return *this; + } + + // Enter here only if (iteration_contiguous_ == + // ThreadMap::Iteration::kContiguous) + the_predicates.iteration_contiguous_ = 0; + ++the_predicates.iteration_strided_; + + if (the_predicates.iteration_strided_ < ThreadMap::Iterations::kStrided) { + pointer_ += params_.inc_strided_; + + return *this; + } + + // Enter here only if (iteration_stride_ == ThreadMap::Iteration::kStrided) + // which means we enter the next tile. + the_predicates.iteration_strided_ = 0; + + // advance to next tile + pointer_ += params_.inc_next_; + + // now return to start tile - if the iterator is subsequently advanced, this + // subtraction as well as the subsequent integer addition are both elided by + // the compiler. + pointer_ -= params_.inc_advance_; + + return *this; + } + + /// Increment and return an instance to self. + CUTLASS_HOST_DEVICE + CustomPredicatedTileAccessIterator operator++(int) { + CustomPredicatedTileAccessIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + the_predicates.clear_mask(enable); + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { + the_predicates.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { + the_predicates.set_mask(mask); + } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { + the_predicates.get_mask(mask); + } + + /// Returns whether access is valid or not + CUTLASS_HOST_DEVICE + bool valid() const { + return the_predicates.valid(); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of CustomPredicatedTileAccessIterator for column-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + int NADim, + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + typename AccessType_> +class CustomPredicatedTileAccessIterator< + NADim, + Shape_, + Element_, + layout::ColumnMajor, + AdvanceRank, + ThreadMap_, + AccessType_> { + public: + static_assert(NADim >= 1 && NADim < 4); + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Dim = typename natten::cuda::fna::GetDim::type; + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingIterator = CustomPredicatedTileAccessIterator< + NADim, + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 0 : 1), + ThreadMap, + AccessType>; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend CustomPredicatedTileAccessIterator; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + /// Default constructor + Params() = default; + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Dim stride, Dim extent_row) : params_(stride, extent_row){}; + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const& base) + : params_(base) {} + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Default constructor + CustomPredicatedTileAccessIterator() = default; + + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + CustomPredicatedTileAccessIterator( + ///< Precomputed parameters object + Params const& params, + ///< Pointer to start of tensor + Pointer pointer, + /// Extent of tensor (rows) + Dim extent_row, + /// Extent of tensor (columns) + int32_t extent_col, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const& threadblock_offset) + : iterator_( + params.params_, + pointer, + extent_row, + extent_col, + thread_id, + layout::PitchLinearCoord( + threadblock_offset.row(), + threadblock_offset.column())) {} + + /// Construct a CustomPredicatedTileAccessIterator with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + CustomPredicatedTileAccessIterator( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + Dim extent_row, ///< Extent of tensor (rows) + int32_t extent_col, ///< Extent of tensor (columns) + int thread_id ///< ID of each participating thread + ) + : CustomPredicatedTileAccessIterator( + params, + pointer, + extent_row, + extent_col, + thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) { + iterator_.set_iteration_index(index); + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const& tile_offset) { + iterator_.add_tile_offset({tile_offset.row(), tile_offset.column()}); + } + + /// Returns a pointer + CUTLASS_HOST_DEVICE + AccessType* get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + CustomPredicatedTileAccessIterator& operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + CustomPredicatedTileAccessIterator operator++(int) { + CustomPredicatedTileAccessIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + iterator_.clear_mask(enable); + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { + iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { + iterator_.set_mask(mask); + } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { + iterator_.get_mask(mask); + } + + /// Returns whether access is valid or not + CUTLASS_HOST_DEVICE + bool valid() { + return iterator_.valid(); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of CustomPredicatedTileAccessIterator for row-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + int NADim, + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + typename AccessType_> +class CustomPredicatedTileAccessIterator< + NADim, + Shape_, + Element_, + layout::RowMajor, + AdvanceRank, + ThreadMap_, + AccessType_> { + public: + static_assert(NADim >= 1 && NADim < 4); + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Dim = typename natten::cuda::fna::GetDim::type; + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingIterator = CustomPredicatedTileAccessIterator< + NADim, + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap, + AccessType>; + + static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend CustomPredicatedTileAccessIterator; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + /// Default constructor + Params() = default; + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Dim stride, Dim extent_row) : params_(stride, extent_row){}; + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const& base) + : params_(base) {} + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Default constructor + CustomPredicatedTileAccessIterator() = default; + + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + CustomPredicatedTileAccessIterator( + ///< Precomputed parameters object + Params const& params, + ///< Pointer to start of tensor + Pointer pointer, + /// Extent of tensor (rows) + Dim extent_row, + /// Extent of tensor (columns) + int32_t extent_col, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const& threadblock_offset) + : iterator_( + params.params_, + pointer, + extent_row, + extent_col, + thread_id, + layout::PitchLinearCoord( + threadblock_offset.column(), + threadblock_offset.row())) {} + + /// Construct a CustomPredicatedTileAccessIterator with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + CustomPredicatedTileAccessIterator( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + Dim extent_row, ///< Extent of tensor (rows) + int32_t extent_col, ///< Extent of tensor (columns) + int thread_id ///< ID of each participating thread + ) + : CustomPredicatedTileAccessIterator( + params, + pointer, + extent_row, + extent_col, + thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) { + iterator_.set_iteration_index(index); + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const& tile_offset) { + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + } + + /// Returns a pointer + CUTLASS_HOST_DEVICE + AccessType* get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + CustomPredicatedTileAccessIterator& operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + CustomPredicatedTileAccessIterator operator++(int) { + CustomPredicatedTileAccessIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + iterator_.clear_mask(enable); + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { + iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { + iterator_.set_mask(mask); + } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { + iterator_.get_mask(mask); + } + + /// Returns whether access is valid or not + CUTLASS_HOST_DEVICE + bool valid() { + return iterator_.valid(); + } +}; + +} // namespace threadblock +} // namespace transform +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fna/iterators/predicated_tile_access_iterator_residual_last.h b/natten/csrc/include/natten/cuda/fna/iterators/predicated_tile_access_iterator_residual_last.h new file mode 100644 index 00000000..2634afac --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna/iterators/predicated_tile_access_iterator_residual_last.h @@ -0,0 +1,976 @@ +/* + * Copied from xFormers (https://github.com/facebookresearch/xformers/) and + * edited. + * + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + */ +/*************************************************************************************************** + * Copyright (c) 2017 - 2026 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 Templates calculating the address and predicates to the load of tiles + from pitch-linear rank=2 tensors. + + This iterator uses masks to guard out-of-bounds accesses. The first tile + this iterator visits maybe partial, then the remaining tiles are complete. + So, we only need to compute the predicates twice, once before the first tile + and once for the remaining full tiles which can share the same predicates. + + A precomputed "Params" object minimizes the amount of state that must be + stored in registers, and integer addition is used to advance the pointer + through memory. +*/ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +//////////////////////////////////////////////////////////////////////////////// + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace transform { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// PredicatedTileAccessIteratorResidualLast +/// +template < + int NADim, + typename Shape, + typename Element, + typename Layout, + int AdvanceRank, + typename ThreadMap, + typename AccessType> +class PredicatedTileAccessIteratorResidualLast; + +//////////////////////////////////////////////////////////////////////////////// + +/// PredicatedTileAccessIteratorResidualLast for possibly multiple axes (NA1D, +/// 2D, 3D) +template < + int NADim, + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + typename AccessType_> +class PredicatedTileAccessIteratorResidualLast< + NADim, + Shape_, + Element_, + layout::PitchLinear, + AdvanceRank, + ThreadMap_, + AccessType_> { + public: + static_assert(NADim >= 1 && NADim < 4); + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Dim = typename natten::cuda::fna::GetDim::type; + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::PitchLinear; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingPredicates = CustomPredicatedTileAccessIteratorPredicates< + NADim, + Shape, + Element, + Layout, + AdvanceRank, + ThreadMap, + AccessType>; + + static int const kAccessesPerVector = + ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert( + !(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + using Mask = typename UnderlyingPredicates::Mask; + + /// Uses a non-template class + struct Params : CustomPredicatedTileAccessIteratorParams { + using Base = CustomPredicatedTileAccessIteratorParams; + + // Default ctor + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Dim stride, Dim extent_row) + : Base( + stride, + extent_row, + MakePredicatedTileAccessIteratorDesc< + Shape, + Element, + Layout, + kAdvanceRank, + ThreadMap>()()) {} + + CUTLASS_HOST_DEVICE + Params(Base const& base) : Base(base) {} + }; + + private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char*; + + private: + // + // Data members + // + + UnderlyingPredicates the_predicates; + Mask residual_tile_mask; + + /// Parameters object with precomputed internal state + Params params_; + + /// Internal pointer to first access of tile + BytePointer pointer_; + TensorCoord coord_offset_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor (rows) + Dim extent_row, + /// Extent of tensor (columns) + int32_t extent_col, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const& threadblock_offset) + : params_(params), + pointer_(reinterpret_cast( + const_cast(pointer))), + the_predicates(extent_row, extent_col) { + the_predicates.set_predicates(thread_id, threadblock_offset); + the_predicates.get_mask(residual_tile_mask); + + // Working around a weird compiler bug happening on P100 for the backward. + // I've seen together: the_predicates.predicates_[0] = 14 (instead of 15) + // residual_tile_mask[0] = 15 (correct) + // + // Adding prints when the value is calculated (in `compute_predicates_`) + // sometimes removes the bug. The consequence is that we skip some + // element of a tensor, leading to wrong results + // Setting `compute_predicates_`'s second argument (`is_steady_state`) to + // true also seems to get rid of the bug - at the cost of twice as many + // comparisons. +#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 700) + constexpr bool kWorkAroundCompilerBug = false; +#else + constexpr bool kWorkAroundCompilerBug = true; +#endif + the_predicates.compute_predicates_( + the_predicates.extent_, true && !kWorkAroundCompilerBug); + + // update internal pointers + coord_offset_ = the_predicates.thread_offset_; + // add_pointer_offset(the_predicates.thread_offset_.strided(), + // the_predicates.thread_offset_.contiguous()); + } + + /// Construct a PredicatedTileAccessIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor (rows) + Dim extent_row, + /// Extent of tensor (columns) + int32_t extent_col, + ///< ID of each participating thread + int thread_id) + : PredicatedTileAccessIteratorResidualLast( + params, + pointer, + extent_row, + extent_col, + thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) { + the_predicates.set_iteration_index(index); + } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool is_residual_tile) { + if (is_residual_tile) { + the_predicates.set_mask(residual_tile_mask); + } + } + + CUTLASS_HOST_DEVICE + void add_pointer_offset(int32_t row_offset, int32_t col_offset) { + coord_offset_.strided() += row_offset; + coord_offset_.contiguous() += col_offset; + } + + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex linear_offset) { + printf("FATAL!"); + // if constexpr (kAdvanceRank) { + // auto col_offset = linear_offset % Shape::kStrided; + // auto row_offset = linear_offset / Shape::kStrided; + // add_pointer_offset((int32_t)row_offset, (int32_t)col_offset); + // } else { + // auto col_offset = linear_offset % Shape::kContiguous; + // auto row_offset = linear_offset / Shape::kContiguous; + // add_pointer_offset((int32_t)row_offset, (int32_t)col_offset); + // } + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_DEVICE + void add_tile_offset(TensorCoord const& tile_offset) { + add_pointer_offset( + Shape::kStrided * tile_offset.strided(), + Shape::kContiguous * tile_offset.contiguous()); + // if (kAdvanceRank) { + // add_pointer_offset( + // params_.inc_advance_ * (tile_offset.strided()) + + // Shape::kContiguous * tile_offset.contiguous()); + // // pointer_ += params_.inc_advance_ * LongIndex(tile_offset.strided()); + // // pointer_ += Shape::kContiguous * tile_offset.contiguous(); + // } else { + // // Advance along contiguous dimension + // add_pointer_offset( + // params_.inc_advance_ * (tile_offset.contiguous()) + + // Shape::kStrided * tile_offset.strided()); + // // pointer_ += params_.inc_advance_ * + // LongIndex(tile_offset.contiguous()); + // // pointer_ += Shape::kStrided * tile_offset.strided(); + // } + } + + CUTLASS_HOST_DEVICE + LongIndex get_pointer_offset() const { + LongIndex coord_contig = coord_offset_.contiguous() + + the_predicates.iteration_contiguous_ * ThreadMap::Delta::kContiguous; + LongIndex coord_strided = + (natten::cuda::fna::map_index_to_coord( + coord_offset_.strided() + + the_predicates.iteration_strided_ * ThreadMap::Delta::kStrided, + the_predicates.extent_row_dim) * + params_.stride_) + .sum(); + return ((coord_strided + coord_contig) * sizeof_bits::value) / 8; + } + + /// Returns a pointer + CUTLASS_HOST_DEVICE + AccessType* get() const { + if (!valid()) { + return nullptr; + } + return reinterpret_cast(pointer_ + get_pointer_offset()) + + the_predicates.iteration_vector_; + } + + /// Increment and return an instance to self. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast& operator++() { + the_predicates.operator++(); + + ++the_predicates.iteration_vector_; + if (the_predicates.iteration_vector_ < kAccessesPerVector) { + return *this; + } + + the_predicates.iteration_vector_ = 0; + ++the_predicates.iteration_contiguous_; + + if (the_predicates.iteration_contiguous_ < + ThreadMap::Iterations::kContiguous) { + return *this; + } + + // Enter here only if (iteration_contiguous_ == + // ThreadMap::Iteration::kContiguous) + the_predicates.iteration_contiguous_ = 0; + ++the_predicates.iteration_strided_; + + if (the_predicates.iteration_strided_ < ThreadMap::Iterations::kStrided) { + // add_pointer_offset(params_.inc_strided_); + // pointer_ += params_.inc_strided_; + + return *this; + } + + // Enter here only if (iteration_stride_ == ThreadMap::Iteration::kStrided) + // which means we enter the next tile. + the_predicates.iteration_strided_ = 0; + + //// advance to next tile + // add_pointer_offset(params_.inc_next_); + //// pointer_ += params_.inc_next_; + + // add_pointer_offset(-1 * params_.inc_advance_); + ////if constexpr (kAdvanceRank) { + //// add_pointer_offset(-1 * params_.inc_advance_, 0); + ////} else { + //// // along the contiguous axis. + //// add_pointer_offset(0, -1 * params_.inc_advance_); + ////} + //// now return to start tile - if the iterator is subsequently advanced, + /// this / subtraction as well as the subsequent integer addition are both + /// elided by / the compiler. / pointer_ -= params_.inc_advance_; + + return *this; + } + + /// Increment and return an instance to self. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast operator++(int) { + PredicatedTileAccessIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + the_predicates.clear_mask(enable); + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { + the_predicates.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { + the_predicates.set_mask(mask); + } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { + the_predicates.get_mask(mask); + } + + /// Returns whether access is valid or not + CUTLASS_HOST_DEVICE + bool valid() const { + return the_predicates.valid(); + } +}; + +// Specialization for NA1D, where we only have one stride, +// and therefore the typical pitch-linear layout, and won't have +// to manually convert offsets on every get(). +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + typename AccessType_> +class PredicatedTileAccessIteratorResidualLast< + 1, + Shape_, + Element_, + layout::PitchLinear, + AdvanceRank, + ThreadMap_, + AccessType_> { + public: + static constexpr int NADim = 1; + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Dim = typename natten::cuda::fna::GetDim::type; + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::PitchLinear; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingPredicates = CustomPredicatedTileAccessIteratorPredicates< + NADim, + Shape, + Element, + Layout, + AdvanceRank, + ThreadMap, + AccessType>; + + static int const kAccessesPerVector = + ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert( + !(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + using Mask = typename UnderlyingPredicates::Mask; + + /// Uses a non-template class + struct Params : CustomPredicatedTileAccessIteratorParams { + using Base = CustomPredicatedTileAccessIteratorParams; + + // Default ctor + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Dim stride, Dim extent_row) + : Base( + stride, + extent_row, + MakePredicatedTileAccessIteratorDesc< + Shape, + Element, + Layout, + kAdvanceRank, + ThreadMap>()()) {} + + CUTLASS_HOST_DEVICE + Params(Base const& base) : Base(base) {} + }; + + private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char*; + + private: + // + // Data members + // + + UnderlyingPredicates the_predicates; + Mask residual_tile_mask; + + /// Parameters object with precomputed internal state + Params params_; + + /// Internal pointer to first access of tile + BytePointer pointer_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor (rows) + Dim extent_row, + /// Extent of tensor (columns) + int32_t extent_col, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const& threadblock_offset) + : params_(params), + pointer_(reinterpret_cast( + const_cast(pointer))), + the_predicates(extent_row, extent_col) { + the_predicates.set_predicates(thread_id, threadblock_offset); + the_predicates.get_mask(residual_tile_mask); + + // Working around a weird compiler bug happening on P100 for the backward. + // I've seen together: the_predicates.predicates_[0] = 14 (instead of 15) + // residual_tile_mask[0] = 15 (correct) + // + // Adding prints when the value is calculated (in `compute_predicates_`) + // sometimes removes the bug. The consequence is that we skip some + // element of a tensor, leading to wrong results + // Setting `compute_predicates_`'s second argument (`is_steady_state`) to + // true also seems to get rid of the bug - at the cost of twice as many + // comparisons. +#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 700) + constexpr bool kWorkAroundCompilerBug = false; +#else + constexpr bool kWorkAroundCompilerBug = true; +#endif + the_predicates.compute_predicates_( + the_predicates.extent_, true && !kWorkAroundCompilerBug); + + // update internal pointers + auto offset = (natten::cuda::fna::map_index_to_coord( + the_predicates.thread_offset_.strided(), extent_row) * + params_.stride_) + .sum() + + the_predicates.thread_offset_.contiguous(); + + add_pointer_offset(offset); + } + + /// Construct a PredicatedTileAccessIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor (rows) + Dim extent_row, + /// Extent of tensor (columns) + int32_t extent_col, + ///< ID of each participating thread + int thread_id) + : PredicatedTileAccessIteratorResidualLast( + params, + pointer, + extent_row, + extent_col, + thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) { + the_predicates.set_iteration_index(index); + } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool is_residual_tile) { + if (is_residual_tile) { + the_predicates.set_mask(residual_tile_mask); + } + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += sizeof_bits::value * pointer_offset / 8; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_DEVICE + void add_tile_offset(TensorCoord const& tile_offset) { + if (kAdvanceRank) { + pointer_ += params_.inc_advance_ * LongIndex(tile_offset.strided()); + pointer_ += Shape::kContiguous * tile_offset.contiguous(); + } else { + pointer_ += params_.inc_advance_ * LongIndex(tile_offset.contiguous()); + pointer_ += Shape::kStrided * tile_offset.strided(); + } + } + + /// Returns a pointer + CUTLASS_HOST_DEVICE + AccessType* get() const { + return reinterpret_cast( + pointer_ + + the_predicates.iteration_contiguous_ * + (ThreadMap::Delta::kContiguous * + sizeof_bits::value) / + 8) + + the_predicates.iteration_vector_; + } + + /// Increment and return an instance to self. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast& operator++() { + the_predicates.operator++(); + + ++the_predicates.iteration_vector_; + if (the_predicates.iteration_vector_ < kAccessesPerVector) { + return *this; + } + + the_predicates.iteration_vector_ = 0; + ++the_predicates.iteration_contiguous_; + + if (the_predicates.iteration_contiguous_ < + ThreadMap::Iterations::kContiguous) { + return *this; + } + + // Enter here only if (iteration_contiguous_ == + // ThreadMap::Iteration::kContiguous) + the_predicates.iteration_contiguous_ = 0; + ++the_predicates.iteration_strided_; + + if (the_predicates.iteration_strided_ < ThreadMap::Iterations::kStrided) { + pointer_ += params_.inc_strided_; + + return *this; + } + + // Enter here only if (iteration_stride_ == ThreadMap::Iteration::kStrided) + // which means we enter the next tile. + the_predicates.iteration_strided_ = 0; + + // advance to next tile + pointer_ += params_.inc_next_; + + // now return to start tile - if the iterator is subsequently advanced, + // this subtraction as well as the subsequent integer addition are both + // elided by the compiler. + pointer_ -= params_.inc_advance_; + + return *this; + } + + /// Increment and return an instance to self. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast operator++(int) { + PredicatedTileAccessIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + the_predicates.clear_mask(enable); + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { + the_predicates.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { + the_predicates.set_mask(mask); + } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { + the_predicates.get_mask(mask); + } + + /// Returns whether access is valid or not + CUTLASS_HOST_DEVICE + bool valid() const { + return the_predicates.valid(); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIteratorResidualLast for row-major +/// data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + int NADim, + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + typename AccessType_> +class PredicatedTileAccessIteratorResidualLast< + NADim, + Shape_, + Element_, + layout::RowMajor, + AdvanceRank, + ThreadMap_, + AccessType_> { + public: + static_assert(NADim >= 1 && NADim < 4); + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Dim = typename natten::cuda::fna::GetDim::type; + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingIterator = PredicatedTileAccessIteratorResidualLast< + NADim, + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap, + AccessType>; + + static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileAccessIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + /// Default ctor + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Dim stride, Dim extent_row) : params_(stride, extent_row){}; + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const& base) + : params_(base) {} + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + ///< Precomputed parameters object + Params const& params, + ///< Pointer to start of tensor + Pointer pointer, + /// Extent of tensor (rows) + Dim extent_row, + /// Extent of tensor (columns) + int32_t extent_col, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const& threadblock_offset) + : iterator_( + params.params_, + pointer, + extent_row, + extent_col, + thread_id, + layout::PitchLinearCoord( + threadblock_offset.column(), + threadblock_offset.row())) {} + + /// Construct a PredicatedTileAccessIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + Dim extent_row, ///< Extent of tensor (rows) + int32_t extent_col, ///< Extent of tensor (columns) + int thread_id ///< ID of each participating thread + ) + : PredicatedTileAccessIteratorResidualLast( + params, + pointer, + extent_row, + extent_col, + thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) { + iterator_.set_iteration_index(index); + } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { + iterator_.set_residual_tile(enable); + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const& tile_offset) { + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + } + + /// Returns a pointer + CUTLASS_HOST_DEVICE + AccessType* get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast& operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileAccessIteratorResidualLast operator++(int) { + PredicatedTileAccessIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + iterator_.clear_mask(enable); + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { + iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { + iterator_.set_mask(mask); + } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { + iterator_.get_mask(mask); + } + + /// Returns whether access is valid or not + CUTLASS_HOST_DEVICE + bool valid() { + return iterator_.valid(); + } +}; + +} // namespace threadblock +} // namespace transform +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fna/iterators/predicated_tile_iterator.h b/natten/csrc/include/natten/cuda/fna/iterators/predicated_tile_iterator.h new file mode 100644 index 00000000..a10b1389 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna/iterators/predicated_tile_iterator.h @@ -0,0 +1,831 @@ +/* + * Copied from CUTLASS (https://github.com/NVIDIA/cutlass/) and edited. + */ +/*************************************************************************************************** + * Copyright (c) 2017 - 2026 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 Templates implementing loading of tiles from pitch-linear rank=2 + tensors. + + This iterator uses masks to guard out-of-bounds accesses. The first tile + this iterator visits maybe partial, then the remaining tiles are complete. + So, we only need to compute the predicates twice, once before the first tile + and once for the remaining full tiles which can share the same predicates. + + A precomputed "Params" object minimizes the amount of state that must be + stored in registers, and integer addition is used to advance the pointer + through memory. +*/ + +#pragma once + +#include + +#include +#include + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace transform { +namespace threadblock { + +template < + int NADim, + typename Shape, + typename Element, + typename Layout, + int AdvanceRank, + typename ThreadMap, + int AccessSize = ThreadMap::kElementsPerAccess> +class CustomPredicatedTileIterator; + +//////////////////////////////////////////////////////////////////////////////// + +template < + int NADim, + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int AccessSize> +class CustomPredicatedTileIterator< + NADim, + Shape_, + Element_, + layout::PitchLinear, + AdvanceRank, + ThreadMap_, + AccessSize> { + public: + static_assert(NADim >= 1 && NADim < 4); + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Dim = typename natten::cuda::fna::GetDim::type; + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::PitchLinear; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + /// Type used for internal memory accesses + using AccessType = AlignedArray< + Element, + AccessSize, + (AccessSize * sizeof_bits::value / 8)>; + + /// Underlying iterator to compute the addresses + using TileAccessIterator = CustomPredicatedTileAccessIterator< + NADim, + Shape, + Element, + Layout, + kAdvanceRank, + ThreadMap, + AccessType>; + + static int const kAccessesPerVector = TileAccessIterator::kAccessesPerVector; + + /// Fragment object to be loaded or stored + using Fragment = cutlass::Array< + Element, + ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; + + /// Predicate vector stores mask to guard accesses + using Mask = typename TileAccessIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + public: + using Base = typename TileAccessIterator::Params::Base; + + friend CustomPredicatedTileIterator; + + private: + /// Parameters object + typename TileAccessIterator::Params params_; + + public: + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Dim stride, Dim extent_row) : params_(stride, extent_row) {} + + /// Default constructor + Params() = default; + + CUTLASS_HOST_DEVICE + Params(Base const& base) : params_(base) {} + }; + + private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char*; + + private: + // + // Data members + // + + /// Data member to the tile access iterator + TileAccessIterator address_iterator_; + + public: + /// Default constructor + CustomPredicatedTileIterator() = default; + + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + CustomPredicatedTileIterator( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor (rows) + Dim extent_row, + /// Extent of tensor (columns) + int32_t extent_col, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const& threadblock_offset) + : address_iterator_( + params.params_, + pointer, + extent_row, + extent_col, + thread_id, + threadblock_offset) {} + + /// Construct a CustomPredicatedTileIterator with zero threadblock offset + CUTLASS_HOST_DEVICE + CustomPredicatedTileIterator( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + Dim extent_row, ///< Extent of tensor (rows) + int32_t extent_col, ///< Extent of tensor (columns) + int thread_id ///< ID of each participating thread + ) + : CustomPredicatedTileIterator( + params, + pointer, + extent_row, + extent_col, + thread_id, + make_Coord(0, 0)) {} + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + address_iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + CustomPredicatedTileIterator& operator++() { + if (kAdvanceRank) + address_iterator_.add_tile_offset({0, 1}); + else + address_iterator_.add_tile_offset({1, 0}); + + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + CustomPredicatedTileIterator operator++(int) { + CustomPredicatedTileIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + address_iterator_.clear_mask(enable); + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { + address_iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { + address_iterator_.set_mask(mask); + } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { + address_iterator_.get_mask(mask); + } + + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment& frag, Index pointer_offset) { + load_with_byte_offset( + frag, pointer_offset * sizeof_bits::value / 8); + } + + CUTLASS_DEVICE + void load_with_byte_offset(Fragment& frag, LongIndex byte_offset) { + AccessType* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + int idx = v + + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); + + address_iterator_.set_iteration_index(idx); + char const* byte_ptr = + reinterpret_cast(address_iterator_.get()) + + byte_offset; + + AccessType const* access_ptr = + reinterpret_cast(byte_ptr); + + cutlass::arch::global_load( + frag_ptr[idx], access_ptr, address_iterator_.valid()); + + ++address_iterator_; + } + } + } + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) { + load_with_byte_offset(frag, 0); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) { + store_with_byte_offset( + frag, pointer_offset * sizeof_bits::value / 8); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, LongIndex byte_offset) { + address_iterator_.set_iteration_index(0); + AccessType const* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + int idx = v + + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); + + char* byte_ptr = + reinterpret_cast(address_iterator_.get()) + byte_offset; + AccessType* access_ptr = reinterpret_cast(byte_ptr); + + if (address_iterator_.valid()) { + *access_ptr = frag_ptr[idx]; + } + ++address_iterator_; + } + } + } + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { + store_with_byte_offset(frag, 0); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of CustomPredicatedTileIterator for column-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + int NADim, + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int AccessSize> +class CustomPredicatedTileIterator< + NADim, + Shape_, + Element_, + layout::ColumnMajor, + AdvanceRank, + ThreadMap_, + AccessSize> { + public: + static_assert(NADim >= 1 && NADim < 4); + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Dim = typename natten::cuda::fna::GetDim::type; + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingIterator = CustomPredicatedTileIterator< + NADim, + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 0 : 1), + ThreadMap, + AccessSize>; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = cutlass::Array< + Element, + ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend CustomPredicatedTileIterator; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + /// Default constructor + Params() = default; + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Dim stride, Dim extent_row) : params_(stride, extent_row) {} + + CUTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const& base) + : params_(base) {} + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Default constructor + CustomPredicatedTileIterator() = default; + + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + CustomPredicatedTileIterator( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + Dim extent_row, ///< Extent of tensor (rows) + int32_t extent_col, ///< Extent of tensor (columns) + int thread_id, ///< ID of each participating thread + TensorCoord const& threadblock_offset ///< Initial offset of threadblock + ) + : iterator_( + params.params_, + pointer, + extent_row, + extent_col, + thread_id, + layout::PitchLinearCoord( + threadblock_offset.row(), + threadblock_offset.column())) {} + + /// Construct a CustomPredicatedTileIterator with zero threadblock offset + CUTLASS_HOST_DEVICE + CustomPredicatedTileIterator( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + Dim extent_row, ///< Extent of tensor (rows) + int32_t extent_col, ///< Extent of tensor (columns) + int thread_id ///< ID of each participating thread + ) + : CustomPredicatedTileIterator( + params, + pointer, + extent_row, + extent_col, + thread_id, + make_Coord(0, 0)) {} + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + CustomPredicatedTileIterator& operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + CustomPredicatedTileIterator operator++(int) { + CustomPredicatedTileIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + iterator_.clear_mask(enable); + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { + iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { + iterator_.set_mask(mask); + } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { + iterator_.get_mask(mask); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment& frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_byte_offset(Fragment& frag, LongIndex byte_offset) { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) { + load_with_pointer_offset(frag, 0); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, LongIndex byte_offset) { + iterator_.store_with_byte_offset(frag, byte_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { + store_with_pointer_offset(frag, 0); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of CustomPredicatedTileIterator for row-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + int NADim, + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int AccessSize> +class CustomPredicatedTileIterator< + NADim, + Shape_, + Element_, + layout::RowMajor, + AdvanceRank, + ThreadMap_, + AccessSize> { + public: + static_assert(NADim >= 1 && NADim < 4); + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Dim = typename natten::cuda::fna::GetDim::type; + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingIterator = CustomPredicatedTileIterator< + NADim, + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap, + AccessSize>; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = cutlass::Array< + Element, + ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend CustomPredicatedTileIterator; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + /// Default constructor + Params() = default; + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Dim stride, Dim extent_row) : params_(stride, extent_row) {} + + CUTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const& base) + : params_(base) {} + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Default constructor + CustomPredicatedTileIterator() = default; + + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + CustomPredicatedTileIterator( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + Dim extent_row, ///< Extent of tensor (rows) + int32_t extent_col, ///< Extent of tensor (columns) + int thread_id, ///< ID of each participating thread + TensorCoord const& threadblock_offset ///< Initial offset of threadblock + ) + : iterator_( + params.params_, + pointer, + extent_row, + extent_col, + thread_id, + layout::PitchLinearCoord( + threadblock_offset.column(), + threadblock_offset.row())) {} + + /// Construct a CustomPredicatedTileIterator with zero threadblock offset + CUTLASS_HOST_DEVICE + CustomPredicatedTileIterator( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + Dim extent_row, ///< Extent of tensor (rows) + int32_t extent_col, ///< Extent of tensor (columns) + int thread_id ///< ID of each participating thread + ) + : CustomPredicatedTileIterator( + params, + pointer, + extent_row, + extent_col, + thread_id, + make_Coord(0, 0)) {} + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + CustomPredicatedTileIterator& operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + CustomPredicatedTileIterator operator++(int) { + CustomPredicatedTileIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + iterator_.clear_mask(enable); + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { + iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { + iterator_.set_mask(mask); + } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { + iterator_.get_mask(mask); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment& frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_byte_offset(Fragment& frag, LongIndex byte_offset) { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) { + load_with_pointer_offset(frag, 0); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, LongIndex byte_offset) { + iterator_.store_with_byte_offset(frag, byte_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { + store_with_pointer_offset(frag, 0); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fna/iterators/predicated_tile_iterator_residual_last.h b/natten/csrc/include/natten/cuda/fna/iterators/predicated_tile_iterator_residual_last.h new file mode 100644 index 00000000..f7fba893 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna/iterators/predicated_tile_iterator_residual_last.h @@ -0,0 +1,617 @@ +/* + * Copied from xFormers (https://github.com/facebookresearch/xformers/) and + * edited. + * + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + */ +/*************************************************************************************************** + * Copyright (c) 2017 - 2026 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 Templates implementing loading of tiles from pitch-linear rank=2 + tensors. + + This iterator uses masks to guard out-of-bounds accesses. The first tile + this iterator visits maybe partial, then the remaining tiles are complete. + So, we only need to compute the predicates twice, once before the first tile + and once for the remaining full tiles which can share the same predicates. + + A precomputed "Params" object minimizes the amount of state that must be + stored in registers, and integer addition is used to advance the pointer + through memory. +*/ + +#pragma once + +#include + +#include +#include + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace transform { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +template < + int NADim, + typename Shape, + typename Element, + typename Layout, + int AdvanceRank, + typename ThreadMap, + int AccessSize = ThreadMap::kElementsPerAccess> +class PredicatedTileIteratorResidualLast; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIteratorResidualLast for pitch-linear data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + int NADim, + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int AccessSize> +class PredicatedTileIteratorResidualLast< + NADim, + Shape_, + Element_, + layout::PitchLinear, + AdvanceRank, + ThreadMap_, + AccessSize> { + public: + static_assert(NADim >= 1 && NADim < 4); + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Dim = typename natten::cuda::fna::GetDim::type; + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::PitchLinear; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + /// Type used for internal memory accesses + using AccessType = AlignedArray< + Element, + AccessSize, + (AccessSize * sizeof_bits::value / 8)>; + + /// Underlying iterator to compute the addresses + using TileAccessIterator = PredicatedTileAccessIteratorResidualLast< + NADim, + Shape, + Element, + Layout, + kAdvanceRank, + ThreadMap, + AccessType>; + + static int const kAccessesPerVector = TileAccessIterator::kAccessesPerVector; + + /// Fragment object to be loaded or stored + using Fragment = cutlass::Array< + Element, + ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; + + /// Predicate vector stores mask to guard accesses + using Mask = typename TileAccessIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + public: + using Base = typename TileAccessIterator::Params::Base; + + friend PredicatedTileIteratorResidualLast; + + private: + /// Parameters object + typename TileAccessIterator::Params params_; + + public: + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Dim stride, Dim extent_row) : params_(stride, extent_row) {} + + CUTLASS_HOST_DEVICE + Params() {} + + CUTLASS_HOST_DEVICE + Params(Base const& base) : params_(base) {} + }; + + private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char*; + + private: + // + // Data members + // + + /// Data member to the tile access iterator + TileAccessIterator address_iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor (rows) + Dim extent_row, + /// Extent of tensor (columns) + int32_t extent_col, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const& threadblock_offset) + : address_iterator_( + params.params_, + pointer, + extent_row, + extent_col, + thread_id, + threadblock_offset) {} + + /// Construct a PredicatedTileIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + Dim extent_row, ///< Extent of tensor (rows) + int32_t extent_col, ///< Extent of tensor (columns) + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIteratorResidualLast( + params, + pointer, + extent_row, + extent_col, + thread_id, + make_Coord(0, 0)) {} + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + address_iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast& operator++() { + if (kAdvanceRank) + address_iterator_.add_tile_offset({0, 1}); + else + address_iterator_.add_tile_offset({1, 0}); + + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast operator++(int) { + PredicatedTileIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + address_iterator_.clear_mask(enable); + } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { + address_iterator_.set_residual_tile(enable); + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { + address_iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { + address_iterator_.set_mask(mask); + } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { + address_iterator_.get_mask(mask); + } + + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment& frag, Index pointer_offset) { + load_with_byte_offset( + frag, pointer_offset * sizeof_bits::value / 8); + } + + CUTLASS_DEVICE + void load_with_byte_offset(Fragment& frag, LongIndex byte_offset) { + AccessType* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + int idx = v + + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); + + address_iterator_.set_iteration_index(idx); + char const* byte_ptr = + reinterpret_cast(address_iterator_.get()) + + byte_offset; + + AccessType const* access_ptr = + reinterpret_cast(byte_ptr); + + cutlass::arch::global_load( + frag_ptr[idx], access_ptr, address_iterator_.valid()); + + ++address_iterator_; + } + } + } + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) { + load_with_byte_offset(frag, 0); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) { + store_with_byte_offset( + frag, pointer_offset * sizeof_bits::value / 8); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, LongIndex byte_offset) { + address_iterator_.set_iteration_index(0); + AccessType const* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + int idx = v + + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); + + char* byte_ptr = + reinterpret_cast(address_iterator_.get()) + byte_offset; + AccessType* access_ptr = reinterpret_cast(byte_ptr); + + if (address_iterator_.valid()) { + *access_ptr = frag_ptr[idx]; + } + ++address_iterator_; + } + } + } + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { + store_with_byte_offset(frag, 0); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIteratorResidualLast for row-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + int NADim, + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int AccessSize> +class PredicatedTileIteratorResidualLast< + NADim, + Shape_, + Element_, + layout::RowMajor, + AdvanceRank, + ThreadMap_, + AccessSize> { + public: + static_assert(NADim >= 1 && NADim < 4); + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Dim = typename natten::cuda::fna::GetDim::type; + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingIterator = PredicatedTileIteratorResidualLast< + NADim, + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap, + AccessSize>; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = cutlass::Array< + Element, + ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Dim stride, Dim extent_row) : params_(stride, extent_row) {} + + CUTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const& base) + : params_(base) {} + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + Dim extent_row, ///< Extent of tensor (rows) + int32_t extent_col, ///< Extent of tensor (columns) + int thread_id, ///< ID of each participating thread + TensorCoord const& threadblock_offset ///< Initial offset of threadblock + ) + : iterator_( + params.params_, + pointer, + extent_row, + extent_col, + thread_id, + layout::PitchLinearCoord( + threadblock_offset.column(), + threadblock_offset.row())) {} + + /// Construct a PredicatedTileIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + Dim extent_row, ///< Extent of tensor (rows) + int32_t extent_col, ///< Extent of tensor (columns) + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIteratorResidualLast( + params, + pointer, + extent_row, + extent_col, + thread_id, + make_Coord(0, 0)) {} + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast& operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast operator++(int) { + PredicatedTileIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + iterator_.clear_mask(enable); + } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { + iterator_.set_residual_tile(enable); + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { + iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { + iterator_.set_mask(mask); + } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { + iterator_.get_mask(mask); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment& frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_byte_offset(Fragment& frag, LongIndex byte_offset) { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) { + load_with_pointer_offset(frag, 0); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, LongIndex byte_offset) { + iterator_.store_with_byte_offset(frag, byte_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { + store_with_pointer_offset(frag, 0); + } +}; + +} // namespace threadblock +} // namespace transform +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fna/iterators/transpose_warp_iterator.h b/natten/csrc/include/natten/cuda/fna/iterators/transpose_warp_iterator.h new file mode 100644 index 00000000..7aee245c --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna/iterators/transpose_warp_iterator.h @@ -0,0 +1,60 @@ +/* + * Copied from xFormers (https://github.com/facebookresearch/xformers/) + * + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * BSD 3-Clause License + * + * 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 names of Facebook, Deepmind Technologies, NYU, NEC + * Laboratories America and IDIAP Research Institute 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 OWNER 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 + +template +struct TransposeWarpIterator { + using Iterator = char; + static bool constexpr kSupportsTranspose = false; +}; + +template < + /// Operand identity + cutlass::gemm::Operand Operand, + /// Data type of A elements + typename Element, + typename InstructionShape, + bool kTranspose> +struct TransposeWarpIterator< + cutlass::gemm::warp:: + WarpIteratorFromSmem> { + using Iterator = cutlass::gemm::warp:: + WarpIteratorFromSmem; + static bool constexpr kSupportsTranspose = true; +}; diff --git a/natten/csrc/include/natten/cuda/fna/iterators/warp_iterator_from_smem.h b/natten/csrc/include/natten/cuda/fna/iterators/warp_iterator_from_smem.h new file mode 100644 index 00000000..989148d0 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna/iterators/warp_iterator_from_smem.h @@ -0,0 +1,290 @@ +/* + * Copied from xFormers (https://github.com/facebookresearch/xformers/). + * + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + */ +/*************************************************************************************************** + * Copyright (c) 2017 - 2026 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 Inspired from + "cutlass/gemm/warp/mma_tensor_op_tile_access_iterator.h" Loads tiles of GEMM + operands from a RowMajor shared-memory layout into registers to use by A100 + TensorCores. + + The difference with "mma_tensor_op_tile_access_iterator.h" is that: + (1) We use "ldmatrix" to load tiles, rather than manual loads (slightly + faster) (2) We support to transpose the operand (eg read `A.transpose()` when + the shared memory holds `A`) + + This is only implemented for the specific shapes. +*/ +#pragma once + +#include + +//////////////////////////////////////////////////////////////////////////////// +namespace cutlass { +namespace gemm { +namespace warp { + +template < + /// Operand identity + Operand Operand_, + /// Data type of A elements + typename Element_, + typename InstructionShape_, + bool kTranspose = false> +class WarpIteratorFromSmem { + public: + /// Shape of tile to load (concept: MatrixShape) + using Shape = cutlass::MatrixShape<32, 32>; + + /// Operand tag + static Operand const kOperand = Operand_; + static_assert( + kOperand == Operand::kA, + "No support for OperandB at the moment"); + + /// Basic check + static_assert( + kOperand == Operand::kA || kOperand == Operand::kB, + "WarpIteratorFromSmem may only be instantiated for A or B operands to warp-level Mma."); + + /// Element type + using Element = Element_; + static_assert(sizeof_bits::value == 16, "Only supported for half"); + + /// Layout of source tile + using Layout = cutlass::layout::RowMajor; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + static_assert(InstructionShape::kRow == 16, "Only supports 16x8x8 / 16x8x16"); + static_assert( + InstructionShape::kColumn == 8 || InstructionShape::kColumn == 16, + "Only supports 16x8x8 / 16x8x16"); + + /// Delta between *MMA operations (in units of *MMA operations, concept: + /// MatrixShape) + static int const kOpDelta = 1; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Number of elements accessed per Shared Memory load + static int const kElementsPerAccess = + (sizeof_bits::value >= 32 ? 1 + : 32 / sizeof_bits::value); + + using InstructionCount = MatrixShape< + Shape::kRow / InstructionShape::kRow, + Shape::kColumn / InstructionShape::kColumn>; + + static int const kIterations = (kOperand == Operand::kA) + ? InstructionCount::kColumn + : InstructionCount::kRow; + + public: + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = Array< + Element, + (kOperand == Operand::kA) + ? (Shape::kRow* InstructionShape::kColumn / kThreads) + : (Shape::kColumn* InstructionShape::kRow / kThreads)>; + + /// Memory access type + // using AccessType = AlignedArray; + using AccessType = Array; + + static int constexpr kWarpShapeDivisibleInner = + (kOperand == Operand::kA ? InstructionShape::kColumn + : InstructionShape::kRow); + static int constexpr kAccessesInner = + (kWarpShapeDivisibleInner / kElementsPerAccess) / 4; + // Number of 32bits tiles to load per `ldmatrix` + static int const kTilesPerInstruction = InstructionShape::kRow / 8; + static_assert(kTilesPerInstruction == 2, "Only supports 16x8x16 and 16x8x8"); + + private: + /// Underlying tensor reference + TensorRef ref_; + + /// Origin + MatrixCoord origin_; + + /// Iterations in a tile + int iterations_; + + public: + /// Constructor from TensorRef + CUTLASS_HOST_DEVICE + WarpIteratorFromSmem(TensorRef const& ref, int lane_id) + : WarpIteratorFromSmem(ref, {Shape::kRow, Shape::kColumn}, lane_id) {} + CUTLASS_HOST_DEVICE + WarpIteratorFromSmem(TensorRef const& ref, TensorCoord extent, int lane_id) + : ref_(ref), iterations_(0) { + // See also: + // https://docs.nvidia.com/cuda/archive/11.7.1/parallel-thread-execution/index.html#warp-level-matrix-fragment-mma-1688 + // 16x8x8: kAccessesInner = 1 (1 ldmatrix.x4) + // 16x8x16: kAccessesInner = 2 (2 ldmatrix.x4) + int ldsm_vec_num = (lane_id >> 3); + if (kOperand == Operand::kA) { + origin_ = MatrixCoord(lane_id % 8, 0); + static_assert( + InstructionCount::kRow * kTilesPerInstruction == 4, + "can't use ldmatrix.x4"); + int access_m_idx = ldsm_vec_num % kTilesPerInstruction; + int inner_idx = (ldsm_vec_num / kTilesPerInstruction) % kAccessesInner; + int inst_m_idx = ldsm_vec_num / (kTilesPerInstruction * kAccessesInner); + MatrixCoord offset( + access_m_idx * 8 + inst_m_idx * InstructionShape::kRow, + inner_idx * 4 * kElementsPerAccess); + if (kTranspose) { + offset = MatrixCoord(offset.column(), offset.row()); + } + origin_ += offset; + } else { + // XXX: This is not tested or used + origin_ = MatrixCoord(0, lane_id % 8); + static_assert(InstructionCount::kColumn * kAccessesInner == 4, ""); + CUTLASS_PRAGMA_UNROLL + for (int inst_n_idx = 0; inst_n_idx < InstructionCount::kColumn; + ++inst_n_idx) { + CUTLASS_PRAGMA_UNROLL + for (int inner_idx = 0; inner_idx < kAccessesInner; ++inner_idx) { + int access_idx = inner_idx + kAccessesInner * inst_n_idx; + + MatrixCoord offset( + inner_idx * 4 * kElementsPerAccess, inst_n_idx * 8); + + if (access_idx == ldsm_vec_num) { + if (kTranspose) { + offset = MatrixCoord(offset.column(), offset.row()); + } + origin_ += offset; + } + } + } + } + + ref_.add_coord_offset(origin_); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_HOST_DEVICE + WarpIteratorFromSmem& add_tile_offset(TensorCoord const& tile_offset) { + TensorCoord coord_offset( + tile_offset.row() * Shape::kRow, tile_offset.column() * Shape::kColumn); + if (kTranspose) { + coord_offset = TensorCoord{coord_offset.column(), coord_offset.row()}; + } + origin_ += coord_offset; + + ref_.add_coord_offset(coord_offset); + + return *this; + } + + /// Advances the iterator along the advance dimension + CUTLASS_DEVICE + void advance() { + if (kOperand == Operand::kA) { + add_tile_offset({0, 1}); + } else { + add_tile_offset({1, 0}); + } + + iterations_ = 0; + } + + /// increase iterations in a tile + CUTLASS_HOST_DEVICE + WarpIteratorFromSmem& operator++() { + iterations_++; + + if (iterations_ >= kIterations) + advance(); + + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + CUTLASS_DEVICE + void load(Fragment& frag) const { + AccessType* access_ptr = reinterpret_cast(&frag); + using LoadLayout = typename platform:: + conditional::type; + + CUTLASS_PRAGMA_UNROLL + for (int access_m_idx = 0; access_m_idx < + (InstructionCount::kRow * kTilesPerInstruction * kAccessesInner) / 4; + ++access_m_idx) { + MatrixCoord offset; + if (kOperand == Operand::kA) { + offset = MatrixCoord( + access_m_idx * 16, iterations_ * InstructionShape::kColumn); + } else { + offset = MatrixCoord(iterations_ * InstructionShape::kRow, 0); + } + if (kTranspose) { + offset = MatrixCoord(offset.column(), offset.row()); + } + cutlass::arch::ldsm( + access_ptr[access_m_idx], ref_.data() + ref_.offset(offset)); + } + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace cutlass +//////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fna/kernel_backward.h b/natten/csrc/include/natten/cuda/fna/kernel_backward.h new file mode 100644 index 00000000..a2bc53e8 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna/kernel_backward.h @@ -0,0 +1,2763 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Fused Neighborhood Attention kernels are heavily based on the + * memory-efficient attention kernels from the xFormers project by Meta + * Platforms, Inc. + * + * Copyright (c) Facebook, Inc. and its affiliates + * + * BSD 3-Clause License + * + * 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 names of Facebook, Deepmind Technologies, NYU, NEC + * Laboratories America and IDIAP Research Institute 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 OWNER 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 + +// #include +// #include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include + +#include + +#include + +namespace natten { +namespace cuda { +namespace fna { + +namespace { + +template +struct GmemTile { + /* + Helper functions to efficient store/load RF to gmem + + GEMM accumulators have a particular format on A100, and + it takes some compute/shared-memory to rearrange them to + a RowMajor or ColumnMajor format in global memory through + an Epilogue. The same complexity goes for loading into RF. + + This class loads/stores RF as they are, and can be used for + efficient accumulation across gemms for instance: + + ``` + GmemTile tile; + for (int i = 0; i < N; ++i) { + // ... + + Fragment accum; + if (i == 0) { + accum.clear(); + } else { + tile.load(accum); + } + mma(accum, ...); + if (i < N-1) { + // Store for next GEMM + tile.store(accum); + } else { + // Store in tensor (eg RowMajor) + epilogue(accum); + } + + // ... + } + ``` + */ + + // 128bits per thread + using AccessType = cutlass::Array; + static constexpr int32_t kBytes = sizeof(AccessType); + static constexpr int32_t kStride = kNumThreads * AccessType::kElements; + static constexpr int32_t kNumIters = + FragmentType::kElements / AccessType::kElements; + static constexpr int32_t kElementsStored = + kNumThreads * FragmentType::kElements; + static_assert( + FragmentType::kElements % AccessType::kElements == 0, + "fragment not aligned on 128 bits"); + + float* ptr; + + CUTLASS_DEVICE void load(FragmentType& fragment, int thread_id) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kNumIters; ++i) { + AccessType* __restrict__ gmem_ptr = reinterpret_cast( + ptr + thread_id * AccessType::kElements + i * kStride); + AccessType sub_fragment; + cutlass::arch::global_load( + sub_fragment, gmem_ptr, true); + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < AccessType::kElements; ++j) { + fragment[i * AccessType::kElements + j] = sub_fragment[j]; + } + } + } + + CUTLASS_DEVICE void store(FragmentType const& fragment, int thread_id) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kNumIters; ++i) { + AccessType* __restrict__ gmem_ptr = reinterpret_cast( + ptr + thread_id * AccessType::kElements + i * kStride); + AccessType sub_fragment; + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < AccessType::kElements; ++j) { + sub_fragment[j] = fragment[i * AccessType::kElements + j]; + } + cutlass::arch::global_store( + sub_fragment, gmem_ptr, true); + } + } +}; + +struct AtomicLock { + CUTLASS_DEVICE static void acquire( + int32_t* lock, + int set_val, + int thread_id) { + if (thread_id == 0) { + while (atomicCAS(lock, 0 /*cmp*/, set_val /*setval*/) != set_val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 + __nanosleep(40); +#endif + } + } + __syncthreads(); + } + CUTLASS_DEVICE static void release(int32_t* lock, int thread_id) { + if (thread_id == 0) { + int status = 0; +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 + asm volatile("st.global.release.gpu.b32 [%0], %1;\n" + : + : "l"(lock), "r"(status)); +#else + asm volatile("st.global.cg.b32 [%0], %1;\n" : : "l"(lock), "r"(status)); +#endif + } + } +}; + +template +constexpr int getWarpsPerSmBw() { + bool is_half = !cutlass::platform::is_same::value; + if (Arch::kMinComputeCapability >= 80) { + return is_half ? 12 : 8; + } + return 8; +} +} // namespace + +template < + // NA spatial dimension; can be 1 (NA1d), 2 (NA2d), or 3 (NA3d). + int NADim_, + typename CausalMask, + // input/output type + typename scalar_t_, + // Architecture we are targeting (eg `cutlass::arch::Sm80`) + typename ArchTag_, + // run optimized kernel because memory accesses will be aligned + bool kIsAligned_, + // block dimensions + int kBlockSizeI_, + int kBlockSizeJ_, + // upperbound on `max(value.shape[-1], query.shape[-1])` + int kMaxK_ = (int)cutlass::platform::numeric_limits::max(), + // NOTE(alih): we force alignment on dim in NATTEN, so we + // only need kPreload_=true instantiations. + //// when doing a GEMM, preload the next one (uses more shmem) + //// bool kPreload_ = true, + bool kAllowDeltaCompute = false> +struct FusedNeighborhoodAttentionBackwardKernel { + static constexpr int NADim = NADim_; + static_assert(NADim >= 1 && NADim < 4, "Only 1D-3D NA are implemented."); + using Dim = typename GetDim::type; + using NAMask = NeighborhoodAttentionMask; + static constexpr bool kHasCausalDims = CausalMask::AnyCausalDims; + + using scalar_t = scalar_t_; + using output_t = scalar_t; + using output_accum_t = float; + using lse_scalar_t = float; + using accum_t = float; + using ArchTag = ArchTag_; + static constexpr bool kIsAligned = kIsAligned_; + // static constexpr bool kPreload = kPreload_; + static constexpr int kBlockSizeI = kBlockSizeI_; + static constexpr int kBlockSizeJ = kBlockSizeJ_; + static constexpr int kMaxK = kMaxK_; + + static constexpr int64_t kWarpSize = 32; + + // If this is true, we store and accumulate dK/dV in RF + // rather than going back to gmem everytime + static constexpr bool kIsHalf = cutlass::sizeof_bits::value <= 16; + static constexpr bool kOutputInRF = kIsHalf && kMaxK <= kBlockSizeI; + + static constexpr bool kPreload = + kIsHalf && ArchTag::kMinComputeCapability >= 80 && kOutputInRF; + + static_assert( + !kPreload || + (kIsHalf && ArchTag::kMinComputeCapability >= 80 && kOutputInRF), + "preload MMA not supported"); + static constexpr bool kPrologueQK = kPreload; + static constexpr bool kPrologueGV = kPreload; + static constexpr bool kPrologueDOV = kPreload; + static constexpr bool kPrologueGQ = kPreload; + static constexpr bool kPrologueGK = kPreload; + + static constexpr int64_t kNumWarpsPerBlock = + (kBlockSizeI * kBlockSizeJ) / (32 * 32); + + // NOTE(alih): we're just disabling the compute delta fusion. + // Barely got it working; and it looked like it would be slower + // than the unfused reduction kernel in higher rank (the q indexing + // isn't cheap, especially given that each CTA does the full reduction. + + // Compute delta for the f16 kernels + // TODO: Figure out why it's slower on the f32 kernels + // (something due to RF pressure?) + // TODO: Remove condition on `kOutputInRF` - this is needed to work + // around a compiler bug on V100, not exactly sure why but I spent + // too much time on this already. Reproducible with + // (B, Mq, Mkv, K) = (1, 1, 1, 136) for instance + static constexpr bool kKernelComputesDelta = kAllowDeltaCompute && kIsHalf && + (kOutputInRF || ArchTag::kMinComputeCapability != 70); + + // TODO(alih): remove this when it's supported + static_assert( + !kKernelComputesDelta, + "FNA-backwards doesn't support kAllowDeltaCompute yet. Turn it off!"); + + // Launch bounds + static constexpr int64_t kNumThreads = kWarpSize * kNumWarpsPerBlock; + static constexpr int64_t kMinBlocksPerSm = + getWarpsPerSmBw() / kNumWarpsPerBlock; + + using GemmType = gemm_kernel_utils::DefaultGemmType; + using DefaultConfig = + typename cutlass::gemm::device::DefaultGemmConfiguration< + typename GemmType::OpClass, + ArchTag, + scalar_t, + scalar_t, + scalar_t, // ElementC + accum_t // ElementAccumulator + >; + static constexpr auto kOptimalAlignement = cutlass::platform::max( + DefaultConfig::kAlignmentA, + DefaultConfig::kAlignmentB); + static constexpr auto kMinimumAlignment = GemmType::kMinimumAlignment; + + struct MatmulQK { + /* + attn_T = k_j @ q_i.transpose(-2, -1) # matmul + attn_T = (attn_T - logsumexp[i_start:i_end].unsqueeze(1).transpose(-2, + -1)).exp() # epilogue + + with attn_T.shape = (kBlockSizeJ, kBlockSizeI) + */ + using ThreadblockShape = + cutlass::gemm::GemmShape; + using WarpShape = cutlass::gemm::GemmShape<32, 32, GemmType::WarpK>; + + // NOTE(alih): I don't know why xFormers doesn't use FindDefaultMma here. + using DefaultMma_ = typename cutlass::gemm::threadblock::DefaultMma< + scalar_t, // ElementA + cutlass::layout::RowMajor, // LayoutA + kIsAligned ? DefaultConfig::kAlignmentA : GemmType::kMinimumAlignment, + scalar_t, // ElementB + cutlass::layout::ColumnMajor, // LayoutB + kIsAligned ? DefaultConfig::kAlignmentB : GemmType::kMinimumAlignment, + accum_t, // ElementC + cutlass::layout::RowMajor, // LayoutC + typename GemmType::OpClass, + ArchTag, + ThreadblockShape, + WarpShape, + typename GemmType::InstructionShape, + DefaultConfig::kStages, + typename GemmType::Operator, + false, // AccumulatorsInRowMajor = false, + cutlass::gemm::SharedMemoryClearOption::kNone>; + + struct DefaultMma : DefaultMma_ { + using IteratorA_ = typename DefaultMma_::IteratorA; + using IteratorB_ = typename DefaultMma_::IteratorB; + using ThreadblockMma_ = typename DefaultMma_::ThreadblockMma; + + using IteratorA = typename natten::cuda::fna:: + ConvertIterators::NewIteratorA; + using IteratorB = typename natten::cuda::fna:: + ConvertIterators::NewIteratorB; + + using ThreadblockMma = typename natten::cuda::fna:: + ConvertIterators::Mma; + }; + + // using MmaCore = typename DefaultMma::MmaCore; + using Mma = + typename MakeCustomMma::Mma; + + // Epilogue to store to shared-memory in a format that we can use later for + // the second matmul + using B2bGemm = typename cutlass::gemm::threadblock::B2bGemm< + NADim, + typename Mma::Operator::IteratorC, + typename Mma::Operator, + scalar_t, + WarpShape, + ThreadblockShape>; + using AccumLambdaIterator = typename DefaultMmaAccumLambdaIterator< + typename Mma::Operator::IteratorC, + accum_t, + kWarpSize>::Iterator; + using AccumulatorSharedStorage = typename B2bGemm::AccumulatorSharedStorage; + }; + + struct MatmulGradV { + /* + grad_v[j_start:j_end] += attn_T @ do_i # matmul + + Dimensions: (kBlockSizeJ * kNumWarpsPerBlock, kBlockSizeI, K) + (we might need to iterate multiple times on K) + */ + using ThreadblockShape = + cutlass::gemm::GemmShape; + using WarpShape = cutlass::gemm::GemmShape<32, 32, GemmType::WarpK>; + using InstructionShape = typename GemmType::InstructionShape; + + using DefaultGemm = cutlass::gemm::kernel::DefaultGemm< + scalar_t, // ElementA, + cutlass::layout::RowMajor, // LayoutA, + DefaultConfig::kAlignmentA, + scalar_t, // ElementB, + cutlass::layout::RowMajor, // LayoutB, + kIsAligned ? DefaultConfig::kAlignmentB : GemmType::kMinimumAlignment, + output_t, + cutlass::layout::RowMajor, // LayoutC, + accum_t, + typename GemmType::OpClass, + ArchTag, + ThreadblockShape, + WarpShape, + typename GemmType::InstructionShape, + typename DefaultConfig::EpilogueOutputOp, + void, // ThreadblockSwizzle - not used + DefaultConfig::kStages, + false, // SplitKSerial + typename GemmType::Operator>; + + using WarpIteratorA = typename cutlass::gemm::threadblock:: + DefaultWarpIteratorAFromSharedMemory< + typename DefaultGemm::Mma::Operator::Shape, // WarpShape + typename DefaultGemm::Mma::Operator:: + InstructionShape, // InstructionShape + typename DefaultGemm::Mma::Operator:: + IteratorA, // RegularWarpIterator + typename DefaultGemm::Mma::Policy // Policy + >::WarpIterator; + using DefaultMmaFromSmem = + typename cutlass::gemm::threadblock::DefaultMmaFromSharedMemory< + NADim, + typename DefaultGemm::Mma, + MatmulQK::AccumulatorSharedStorage::Shape::kN, + WarpIteratorA, + /*kApplyDropout*/ false>; // kScaleOperandA + + using Mma = typename DefaultMmaFromSmem::Mma; + using IteratorB = typename Mma::IteratorB; + using WarpCount = typename Mma::WarpCount; + + // Epilogue + using DefaultOutputOp = typename DefaultConfig::EpilogueOutputOp; + using DefaultEpilogue = typename DefaultGemm::Epilogue; + using OutputTileIterator = + typename cutlass::epilogue::threadblock::MakePrefetchableIterator< + NADim, + typename DefaultEpilogue::OutputTileIterator>::Iterator; + using AccumTileGmem = GmemTile; + }; + + struct MatmulDOIVJ { + /* + doi_t_vj = do_i @ v_j.transpose(-2, -1) # matmul + tmp = (doi_t_vj - Di.unsqueeze(1)) * attn # inplace / epilogue? + */ + using ThreadblockShape = + cutlass::gemm::GemmShape; + using WarpShape = cutlass::gemm::GemmShape<32, 32, GemmType::WarpK>; + + using ElementC = output_t; + using ElementAccum = accum_t; + + // no-op output op - epilogue just stores result to global memory + using BiasGradEpilogueOutputOp = + typename cutlass::epilogue::thread::LinearCombination< + ElementC, + DefaultConfig::EpilogueOutputOp::kCount, + typename DefaultConfig::EpilogueOutputOp::ElementAccumulator, + typename DefaultConfig::EpilogueOutputOp::ElementCompute, + cutlass::epilogue::thread::ScaleType::Nothing>; + + using DefaultGemm = typename cutlass::gemm::kernel::DefaultGemm< + scalar_t, // ElementA + cutlass::layout::RowMajor, // LayoutA + kIsAligned ? DefaultConfig::kAlignmentA : GemmType::kMinimumAlignment, + scalar_t, // ElementB + cutlass::layout::ColumnMajor, // LayoutB + kIsAligned ? DefaultConfig::kAlignmentB : GemmType::kMinimumAlignment, + ElementC, // ElementC + cutlass::layout::RowMajor, // LayoutC + ElementAccum, // ElementAccumulator + typename GemmType::OpClass, + ArchTag, + ThreadblockShape, + WarpShape, + typename GemmType::InstructionShape, + BiasGradEpilogueOutputOp, // EpilogueOutputOp + void, // ThreadblockSwizzle (not used) + // multiple preloads, and 3 stages push us over shared + // memory capacity on A100. + kPreload && (kBlockSizeI * kBlockSizeJ > 64 * 64) + ? cutlass::const_min(2, DefaultConfig::kStages) + : DefaultConfig::kStages, // Stages + false, // SplitKSerial + typename GemmType::Operator, + cutlass::gemm::SharedMemoryClearOption::kNone>; + using Mma = typename MakeCustomMmaAndReplaceIterators< + NADim, + typename DefaultGemm::Mma, + kMaxK>::Mma; + using AccumLambdaIterator = typename DefaultMmaAccumLambdaIterator< + typename Mma::Operator::IteratorC, + ElementAccum, + kWarpSize>::Iterator; + + // epilogue used to write bias gradient, which is just the output of this + // matmul with some operations applied to the fragment + using BiasGradEpilogue = typename DefaultGemm::Epilogue; + + // Epilogue to store to shared-memory in a format that we can use later for + // the second matmul + using B2bGemm = typename cutlass::gemm::threadblock::B2bGemm< + NADim, + typename DefaultGemm::Mma::Operator::IteratorC, + typename DefaultGemm::Mma::Operator, + scalar_t, + WarpShape, + ThreadblockShape>; + using AccumulatorSharedStorage = typename B2bGemm::AccumulatorSharedStorage; + }; + + struct MatmulGradQ { + // grad_q <- tmp @ k_j + using ThreadblockShape = + cutlass::gemm::GemmShape; + using WarpShape = cutlass::gemm::GemmShape<32, 32, GemmType::WarpK>; + using InstructionShape = typename GemmType::InstructionShape; + + using DefaultGemm = cutlass::gemm::kernel::DefaultGemm< + scalar_t, // ElementA, + cutlass::layout::RowMajor, // LayoutA, + DefaultConfig::kAlignmentA, + scalar_t, // ElementB, + cutlass::layout::RowMajor, // LayoutB, + kIsAligned ? DefaultConfig::kAlignmentB : GemmType::kMinimumAlignment, + output_t, + cutlass::layout::RowMajor, // LayoutC, + accum_t, + typename GemmType::OpClass, + ArchTag, + ThreadblockShape, + WarpShape, + typename GemmType::InstructionShape, + typename DefaultConfig::EpilogueOutputOp, + void, // ThreadblockSwizzle - not used + DefaultConfig::kStages, + false, // SplitKSerial + typename GemmType::Operator>; + + using WarpIteratorA = typename cutlass::gemm::threadblock:: + DefaultWarpIteratorAFromSharedMemory< + typename DefaultGemm::Mma::Operator::Shape, + typename DefaultGemm::Mma::Operator::InstructionShape, + typename DefaultGemm::Mma::Operator::IteratorA, + typename DefaultGemm::Mma::Policy>::WarpIterator; + using DefaultMmaFromSmem = + typename cutlass::gemm::threadblock::DefaultMmaFromSharedMemory< + NADim, + typename DefaultGemm::Mma, + MatmulDOIVJ::AccumulatorSharedStorage::Shape::kN, + WarpIteratorA, + false>; // kScaleOperandA + using Mma = typename DefaultMmaFromSmem::Mma; + using IteratorB = typename Mma::IteratorB; + using WarpCount = typename Mma::WarpCount; + + // Epilogue + using DefaultOutputOp = typename DefaultConfig::EpilogueOutputOp; + using DefaultEpilogue = typename DefaultGemm::Epilogue; + using OutputTileIterator = + typename cutlass::epilogue::threadblock::MakePrefetchableIterator< + NADim, + typename DefaultEpilogue::OutputTileIterator>::Iterator; + using AccumTileGmem = GmemTile; + }; + + struct MatmulGradK { + // grad_k <- tmp.transpose(-2, -1) @ q_i + using ThreadblockShape = + cutlass::gemm::GemmShape; + using WarpShape = cutlass::gemm::GemmShape<32, 32, GemmType::WarpK>; + using InstructionShape = typename GemmType::InstructionShape; + + using DefaultGemm = cutlass::gemm::kernel::DefaultGemm< + scalar_t, // ElementA, + cutlass::layout::RowMajor, // LayoutA, + DefaultConfig::kAlignmentA, + scalar_t, // ElementB, + cutlass::layout::RowMajor, // LayoutB, + kIsAligned ? DefaultConfig::kAlignmentB : GemmType::kMinimumAlignment, + output_t, + cutlass::layout::RowMajor, // LayoutC, + accum_t, + typename GemmType::OpClass, + ArchTag, + ThreadblockShape, + WarpShape, + typename GemmType::InstructionShape, + typename DefaultConfig::EpilogueOutputOp, + void, // ThreadblockSwizzle - not used + DefaultConfig::kStages, + false, // SplitKSerial + typename GemmType::Operator>; + + using WarpIteratorA = typename cutlass::gemm::threadblock:: + DefaultWarpIteratorAFromSharedMemory< + typename DefaultGemm::Mma::Operator::Shape, + typename DefaultGemm::Mma::Operator::InstructionShape, + typename DefaultGemm::Mma::Operator::IteratorA, + typename DefaultGemm::Mma::Policy>::WarpIterator; + using DefaultMmaFromSmemN = + typename cutlass::gemm::threadblock::DefaultMmaFromSharedMemory< + NADim, + typename DefaultGemm::Mma, + MatmulQK::AccumulatorSharedStorage::Shape::kN, // kMaxK + WarpIteratorA, + false>; // kScaleOperandA + using DefaultMmaFromSmemT = + typename cutlass::gemm::threadblock::DefaultMmaFromSharedMemory< + NADim, + typename DefaultGemm::Mma, + MatmulDOIVJ::AccumulatorSharedStorage::Shape::kM, // kMaxK + WarpIteratorA, + false, // kScaleOperandA + kPreload>; // kTransposeA + using DefaultMmaFromSmem = typename cutlass::platform::conditional< + DefaultMmaFromSmemT::kIsTransposedA, + DefaultMmaFromSmemT, + DefaultMmaFromSmemN>::type; + using Mma = typename DefaultMmaFromSmem::Mma; + using IteratorB = typename Mma::IteratorB; + using WarpCount = typename Mma::WarpCount; + + // Epilogue + using DefaultOutputOp = typename DefaultConfig::EpilogueOutputOp; + using DefaultEpilogue = typename DefaultGemm::Epilogue; + using OutputTileIterator = + typename cutlass::epilogue::threadblock::MakePrefetchableIterator< + NADim, + typename DefaultEpilogue::OutputTileIterator>::Iterator; + using AccumTileGmem = GmemTile; + }; + + static constexpr bool kEnableSplitKeys = true; + + static constexpr bool kNeedsAccumGradQ = kEnableSplitKeys || + !cutlass::platform::is_same::value; + static constexpr bool kNeedsAccumGradK = !kOutputInRF && + !cutlass::platform::is_same::value; + static constexpr bool kNeedsAccumGradV = !kOutputInRF && + !cutlass::platform::is_same::value; + + struct GradQTempStorage { + int32_t lock; + int32_t counter; + int32_t pad[2]; // pad to 128bits + output_accum_t buffer[MatmulGradQ::AccumTileGmem::kElementsStored]; + }; + + struct Params { + // Input tensors + scalar_t* query_ptr = nullptr; // [Mq, nH, K] + scalar_t* key_ptr = nullptr; // [Mk, nH, K] + scalar_t* value_ptr = nullptr; // [Mk, nH, Kv] + lse_scalar_t* logsumexp_ptr = nullptr; // [nH, Mq] + scalar_t* output_ptr = nullptr; // [Mq, nH, Kv] + scalar_t* grad_output_ptr = nullptr; // [Mq, nH, Kv] + accum_t* delta_ptr = nullptr; // [nH, Mq] + + // Output tensors + output_t* grad_query_ptr = nullptr; // [Mq, nH, K] + output_t* grad_key_ptr = nullptr; // [Mk, nH, K] + output_t* grad_value_ptr = nullptr; // [Mk, nH, Kv] + + // Accumulators + output_accum_t* workspace = nullptr; // [Mq, Kq] + [Mkv, Kq] + [Mkv, Kv] + output_accum_t* workspace_gv = + nullptr; // (will be calculated by the kernel) + GradQTempStorage* workspace_gq = + nullptr; // (will be calculated by the kernel) + + // Scale + accum_t scale = 1.0f; + + // NA parameters + Dim kernel_size; + Dim stride; + Dim dilation; + + Dim query_tile_shape; + Dim key_tile_shape; + + bool is_fully_block_sparse = false; + bool has_q_padding = false; + + // Dimensions/strides + int32_t head_dim = -1; + int32_t head_dim_value = -1; + + Dim num_queries; + Dim num_queries_post_partitioning; + + int32_t num_heads = -1; + int32_t num_batches = -1; + + Dim lse_strideM; + Dim q_strideM; + Dim k_strideM; + Dim v_strideM; + Dim o_strideM; + + // int16_t num_splits_key = 1; // We use `gridDim.x` inside kernel + Dim num_splits_key; // We use `gridDim.x` inside kernel + + CUTLASS_HOST_DEVICE int32_t num_splits_key_device() const { +#ifdef __CUDA_ARCH__ + return kEnableSplitKeys ? gridDim.x : 1; +#else + return num_splits_key.prod32(); // for host-side tests +#endif + } + + CUTLASS_HOST_DEVICE int16_t split_key_device() const { +#ifdef __CUDA_ARCH__ + return kEnableSplitKeys ? blockIdx.x : 0; +#else + return 0; // for host-side tests +#endif + } + + CUTLASS_DEVICE Dim split_key_device_dim() const { + // auto num_key_tiles = ceil_div_dim(num_queries_post_partitioning, + // key_tile_shape); + return kEnableSplitKeys ? map_index_to_coord(blockIdx.x, num_splits_key) + : Dim(); + } + + CUTLASS_DEVICE bool advance_to_block() { + num_queries_post_partitioning = ceil_div_dim(num_queries, dilation); + // num_tokens = num_queries_post_partitioning.prod32(); + + auto batch_id = blockIdx.z; + auto bidxy = blockIdx.y; + auto dilation_size = dilation.prod32(); + auto dilation_dim_idx = (int32_t)bidxy % dilation_size; + auto dilation_idx = map_index_to_coord(dilation_dim_idx, dilation); + auto head_id = bidxy / dilation_size; + + auto q_strideH = head_dim; + auto k_strideH = head_dim; + auto v_strideH = head_dim_value; + auto o_strideH = head_dim_value; + + auto q_heads_dims = num_heads * q_strideH; + auto k_heads_dims = num_heads * k_strideH; + auto v_heads_dims = num_heads * v_strideH; + auto o_heads_dims = num_heads * o_strideH; + + auto q_stride_dilation = compute_stride(num_queries, q_heads_dims); + auto k_stride_dilation = compute_stride(num_queries, k_heads_dims); + auto v_stride_dilation = compute_stride(num_queries, v_heads_dims); + auto o_stride_dilation = compute_stride(num_queries, o_heads_dims); + + q_strideM = q_stride_dilation * dilation; + k_strideM = k_stride_dilation * dilation; + v_strideM = v_stride_dilation * dilation; + o_strideM = o_stride_dilation * dilation; + + auto q_strideB = num_queries.prod() * q_heads_dims; + auto k_strideB = num_queries.prod() * k_heads_dims; + auto v_strideB = num_queries.prod() * v_heads_dims; + auto o_strideB = num_queries.prod() * o_heads_dims; + + maybe_mask_qk_tiles( + num_queries_post_partitioning, num_queries, dilation, dilation_idx); + + // NOTE: causal mask is antithetical to being fully block sparse. + // While there are cases where a causal mask has mostly dense blocks, the + // last one (at least) will always be causally masked. Fast paths for + // block sparsity in FNA are limited to masks that are guaranteed to be + // fully block sparse only. + is_fully_block_sparse = not kHasCausalDims && + fully_block_sparse( + num_queries_post_partitioning, + kernel_size, + stride, + query_tile_shape, + key_tile_shape); + + has_q_padding = + not evenly_divides(num_queries_post_partitioning, query_tile_shape); + + if constexpr (kNeedsAccumGradQ || kNeedsAccumGradK || kNeedsAccumGradV) { +#ifdef NATTEN_ENABLE_DEVICE_SIDE_ASSERTIONS + assert(workspace_size() == 0 || workspace != nullptr); +#endif + + workspace += (batch_id * num_heads * dilation_size + + head_id * dilation_size + dilation_dim_idx) * + workspace_strideBH(); + // Mutex lock shift + workspace = gemm_kernel_utils::warp_uniform(workspace); + workspace_gv = workspace + workspace_elements_gk(); + workspace_gq = + (GradQTempStorage*)(workspace_gv + workspace_elements_gv()); + + // workspace += dilation_dim_idx * workspace_elements_gk(); + // workspace_gq += dilation_dim_idx * workspace_elements_gq(); + // workspace_gv += dilation_dim_idx * workspace_elements_gv(); + + if (kEnableSplitKeys) { + workspace_gv += workspace_elements_gv() * split_key_device() / + num_splits_key_device(); + workspace += workspace_elements_gk() * split_key_device() / + num_splits_key_device(); + } + } else { + workspace = nullptr; + } + + // Advance pointers that depend on the total concatenated + // number of queries, as `num_queries` is modified in the block + // below + auto lse_stride_dilation = compute_stride(num_queries, num_heads); + lse_strideM = lse_stride_dilation * dilation; + logsumexp_ptr += batch_id * num_queries.prod32() * num_heads + + (dilation_idx * lse_stride_dilation).sum() + head_id; + + // NOTE(alih): assumes delta and lse have identical shape and layout + delta_ptr += batch_id * num_queries.prod32() * num_heads + + (dilation_idx * lse_stride_dilation).sum() + head_id; + + query_ptr += batch_id * q_strideB + head_id * q_strideH + + (dilation_idx * q_stride_dilation).sum(); + key_ptr += batch_id * k_strideB + head_id * k_strideH + + (dilation_idx * k_stride_dilation).sum(); + value_ptr += batch_id * v_strideB + head_id * v_strideH + + (dilation_idx * v_stride_dilation).sum(); + + output_ptr += batch_id * o_strideB + head_id * o_strideH + + (dilation_idx * o_stride_dilation).sum(); + // NOTE(alih): assumes gradient tensors match the layout of the original + // tensors, and are contiguous + // grad_output_ptr += batch_id * gO_strideB + head_id * gO_strideH; + + grad_query_ptr += batch_id * q_strideB + head_id * q_strideH + + (dilation_idx * q_stride_dilation).sum(); + grad_key_ptr += batch_id * k_strideB + head_id * k_strideH + + (dilation_idx * k_stride_dilation).sum(); + grad_value_ptr += batch_id * v_strideB + head_id * v_strideH + + (dilation_idx * v_stride_dilation).sum(); + grad_output_ptr += batch_id * o_strideB + head_id * o_strideH + + (dilation_idx * o_stride_dilation).sum(); + + // Some values are modified above + // Signal to the compiler that they are the same in all threads + // and can be stored in warp-uniform registers (Sm75+) + // num_queries = gemm_kernel_utils::warp_uniform(num_queries); + // num_keys = gemm_kernel_utils::warp_uniform(num_keys); + // custom_mask_type = gemm_kernel_utils::warp_uniform(custom_mask_type); + + query_ptr = gemm_kernel_utils::warp_uniform(query_ptr); + key_ptr = gemm_kernel_utils::warp_uniform(key_ptr); + value_ptr = gemm_kernel_utils::warp_uniform(value_ptr); + // bias_ptr = gemm_kernel_utils::warp_uniform(bias_ptr); + logsumexp_ptr = gemm_kernel_utils::warp_uniform(logsumexp_ptr); + output_ptr = gemm_kernel_utils::warp_uniform(output_ptr); + grad_output_ptr = gemm_kernel_utils::warp_uniform(grad_output_ptr); + delta_ptr = gemm_kernel_utils::warp_uniform(delta_ptr); + + grad_query_ptr = gemm_kernel_utils::warp_uniform(grad_query_ptr); + grad_key_ptr = gemm_kernel_utils::warp_uniform(grad_key_ptr); + grad_value_ptr = gemm_kernel_utils::warp_uniform(grad_value_ptr); + // grad_bias_ptr = gemm_kernel_utils::warp_uniform(grad_bias_ptr); + + return true; + } + + __host__ dim3 getBlocksGrid() const { + return dim3( + num_splits_key.prod32(), num_heads * dilation.prod32(), num_batches); + } + + __host__ dim3 getThreadsGrid() const { + return dim3(kWarpSize * kNumWarpsPerBlock, 1, 1); + } + + CUTLASS_HOST_DEVICE int64_t workspace_elements_gk() const { + if constexpr (!kNeedsAccumGradK) { + return 0; + } + return num_splits_key.prod32() * kBlockSizeJ * + gemm_kernel_utils::align_up(head_dim, (int32_t)kBlockSizeI); + } + + CUTLASS_HOST_DEVICE int64_t workspace_elements_gv() const { + if constexpr (!kNeedsAccumGradV) { + return 0; + } + return num_splits_key.prod32() * kBlockSizeJ * + gemm_kernel_utils::align_up(head_dim_value, (int32_t)kBlockSizeI); + } + + CUTLASS_HOST_DEVICE int64_t workspace_elements_gq() const { + if constexpr (!kNeedsAccumGradQ) { + return 0; + } + auto num_q_post_dilation = ceil_div_dim(num_queries, dilation); + auto num_q_tiles = + ceil_div_dim(num_q_post_dilation, query_tile_shape); // kBlockSizeI + auto num_blocks = num_q_tiles.prod32(); + int num_cols = gemm_kernel_utils::ceil_div( + head_dim, MatmulGradQ::ThreadblockShape::kN); + return num_blocks * num_cols * sizeof(GradQTempStorage) / + sizeof(output_accum_t); + } + + CUTLASS_HOST_DEVICE int64_t workspace_strideBH() const { + // Aligned to 128bits + return gemm_kernel_utils::align_up( + workspace_elements_gk() + workspace_elements_gv() + + workspace_elements_gq(), + int64_t(4)); + } + + CUTLASS_HOST_DEVICE int64_t workspace_size() const { + // Returns size of buffer we need to run this kernel + return (num_batches * num_heads * dilation.prod32() * + workspace_strideBH()) * + sizeof(float); + } + + // NOTE(alih): we check whether a CTA-scope GEMM should load from workspace + // or not, therefore there's no need to zero out the workspace; unless we're + // using split keys, which means the Q gmem tile counters will have to start + // at 0, so we should just wipe everything. + CUTLASS_HOST_DEVICE bool should_zero_workspace() const { + // return num_splits_key > 1 || window_size > 0; + if constexpr (kNeedsAccumGradQ) { + if (num_splits_key.prod32() > 1) { + return true; + } + } + return false; + } + }; + + struct SharedStoragePrologue { + struct { + cutlass::Array di; // (do_i * o_i).sum(-1) + cutlass::Array lse_i; // logsumexp + typename MatmulQK::Mma::SharedStorageA mm_qk_k; + } persistent; + union { + struct { + // part1 - after Q.K / dV / dO.V + union { + // 4. store Pij. it is needed: + // - in dVj += (Pij.T * Zij) @ dOi + // - in dSij = Pij * (dPij - Di) + // 6. dVj += (Pij.T * Zij) @ dOi + // 10. write to fragment + typename MatmulQK::AccumulatorSharedStorage attn_shared_storage; + }; + + union { + // 2. prologue for dVj + // 6. workspace for dVj += (Pij.T * Zij) @ dOi + typename MatmulGradV::Mma::SharedStorage mm_gradV; + // 7. dVj epilogue + typename MatmulGradV::DefaultEpilogue::SharedStorage gradV_epilogue; + }; + + // 3. prologue for dPij_dropped + // 8. used in dPij_dropped = dOi @ Vj.T + typename MatmulDOIVJ::Mma::SharedStorage mm_doivj; + } part1; + + struct { + // part2 - dQ + union { + typename MatmulQK::AccumulatorSharedStorage + tmpT_shared_storage; // (from part1) + typename MatmulDOIVJ::AccumulatorSharedStorage tmp_shared_storage; + }; + typename MatmulGradK::Mma::SharedStorage mm_gradK; // (preload) + typename MatmulGradQ::Mma::SharedStorage mm_gradQ; // (preload) + union { + // store dB = dSij to global memory + typename MatmulDOIVJ::BiasGradEpilogue::SharedStorage gradB_epilogue; + typename MatmulGradQ::DefaultEpilogue::SharedStorage gradQ_epilogue; + }; + + } part2; + + struct { + // part3 - after last iteration on dQ's epilogue / dK + union { + typename MatmulQK::AccumulatorSharedStorage + tmpT_shared_storage; // (from part1) + typename MatmulDOIVJ::AccumulatorSharedStorage tmp_shared_storage; + }; + typename MatmulGradK::Mma::SharedStorage mm_gradK; // (preload) + typename MatmulGradQ::DefaultEpilogue::SharedStorage + gradQ_epilogue_lastIter; + + typename MatmulGradK::DefaultEpilogue::SharedStorage gradK_epilogue; + } part3; + + struct { + // part4 - after last iteration on dK's epilogue / preload next K.Q_t + typename MatmulQK::Mma::SharedStorageB mm_qk_q; + + // If we reach end of current key, dump RF->gmem with "final" epilogues + typename MatmulGradK::DefaultEpilogue::SharedStorage + gradK_epilogue_final; + typename MatmulGradV::DefaultEpilogue::SharedStorage + gradV_epilogue_final; + } part4; + }; + static void print_size() { + // Field size +#define FSZ(f) int((sizeof(((SharedStoragePrologue*)0)->f))) + + printf("Total smem: %d bytes\n", int(sizeof(SharedStoragePrologue))); + printf(" persistent: %db\n", FSZ(persistent)); + printf(" mm_qk_k: %db\n", FSZ(persistent.mm_qk_k)); + printf(" part1: %db\n", FSZ(part1)); + // printf(" bias: %db\n", FSZ(part1.bias)); + printf(" attn_shared_storage: %db\n", FSZ(part1.attn_shared_storage)); + printf(" mm_gradV: %db\n", FSZ(part1.mm_gradV)); + printf(" gradV_epilogue: %db\n", FSZ(part1.gradV_epilogue)); + printf(" mm_doivj: %db\n", FSZ(part1.mm_doivj)); + printf(" part2: %db\n", FSZ(part2)); + printf(" tmpT_shared_storage: %db\n", FSZ(part2.tmpT_shared_storage)); + printf(" tmp_shared_storage: %db\n", FSZ(part2.tmp_shared_storage)); + printf(" mm_gradK: %db\n", FSZ(part2.mm_gradK)); + printf(" mm_gradQ: %db\n", FSZ(part2.mm_gradQ)); + printf(" gradB_epilogue: %db\n", FSZ(part2.gradB_epilogue)); + printf(" gradQ_epilogue: %db\n", FSZ(part2.gradQ_epilogue)); + printf(" part3: %db\n", FSZ(part3)); + printf(" tmpT_shared_storage: %db\n", FSZ(part3.tmpT_shared_storage)); + printf(" part4: %db\n", FSZ(part4)); + printf(" mm_qk_q: %db\n", FSZ(part4.mm_qk_q)); + printf( + " gradK_epilogue_final: %db\n", FSZ(part4.gradK_epilogue_final)); + printf( + " gradV_epilogue_final: %db\n", FSZ(part4.gradV_epilogue_final)); + } +// =========================================== +#define FIELD(INSIDE_STRUCT, FIELDNAME) \ + CUTLASS_DEVICE auto& FIELDNAME() { \ + return INSIDE_STRUCT.FIELDNAME; \ + } + + FIELD(persistent, di) + FIELD(persistent, mm_qk_k) + FIELD(persistent, lse_i) + // FIELD(part1, bias) + FIELD(part1, attn_shared_storage) + FIELD(part1, mm_gradV) + FIELD(part1, gradV_epilogue) + FIELD(part1, mm_doivj) + FIELD(part2, mm_gradK) + FIELD(part2, mm_gradQ) + FIELD(part2, gradB_epilogue) + FIELD(part2, gradQ_epilogue) + FIELD(part2, tmp_shared_storage) + FIELD(part3, tmpT_shared_storage) + FIELD(part3, gradQ_epilogue_lastIter) + FIELD(part3, gradK_epilogue) + FIELD(part4, mm_qk_q) + FIELD(part4, gradK_epilogue_final) + FIELD(part4, gradV_epilogue_final) + }; + + struct SharedStorageNoPrologue { + struct { + cutlass::Array di; // (do_i * o_i).sum(-1) + } persistent; + union { + struct { + // part1 - Q.K matmul + typename MatmulQK::Mma::SharedStorageA mm_qk_k; + typename MatmulQK::Mma::SharedStorageB mm_qk_q; + cutlass::Array lse_i; // logsumexp + } part1; + + struct { + // part2 - compute gradV + union { + // 2. store Pij to shared memory. it is needed: + // - in this step, where it is used in dVj += (Pij.T * Zij) @ dOi + // - in next step where it is used in dSij = Pij * (dPij - Di) + typename MatmulQK::AccumulatorSharedStorage attn_shared_storage; + }; + + union { + typename MatmulGradV::Mma::SharedStorage mm_gradV; + typename MatmulGradV::DefaultEpilogue::SharedStorage gradV_epilogue; + }; + } part2; + + struct { + // part3 - DO.V matmul + union { + // first compute dPij = (dOi @ Vj.T) * Zij + // and dSij = Pij * (dPij - Di) + struct { + // (from part2) - Pij for computing dSij = Pij * (dPij - Di) + typename MatmulQK::AccumulatorSharedStorage attn_shared_storage; + // matmul to compute dOiVj + typename MatmulDOIVJ::Mma::SharedStorage mm_doivj; + }; + // then store dB = dSij to global memory + typename MatmulDOIVJ::BiasGradEpilogue::SharedStorage gradB_epilogue; + }; + } part3; + + struct { + // part4 - compute gradQ + typename MatmulQK::AccumulatorSharedStorage + tmpT_shared_storage; // (from part2) + typename MatmulDOIVJ::AccumulatorSharedStorage tmp_shared_storage; + union { + typename MatmulGradQ::Mma::SharedStorage mm_gradQ; + typename MatmulGradQ::DefaultEpilogue::SharedStorage gradQ_epilogue; + typename MatmulGradQ::DefaultEpilogue::SharedStorage + gradQ_epilogue_lastIter; + }; + } part4; + + struct { + // part5 - compute gradK + typename MatmulQK::AccumulatorSharedStorage + tmpT_shared_storage; // (from part2) + typename MatmulDOIVJ::AccumulatorSharedStorage tmp_shared_storage; + union { + typename MatmulGradK::Mma::SharedStorage mm_gradK; + typename MatmulGradK::DefaultEpilogue::SharedStorage gradK_epilogue; + }; + } part5; + + struct { + // part6 - store RF accumulated into gmem + typename MatmulGradK::DefaultEpilogue::SharedStorage + gradK_epilogue_final; + typename MatmulGradV::DefaultEpilogue::SharedStorage + gradV_epilogue_final; + } part6; + }; + static void print_size() { +#define FIELD_SIZEOF(f) int((sizeof(((SharedStorageNoPrologue*)0)->f))) + printf("Total smem: %d bytes\n", int(sizeof(SharedStorageNoPrologue))); + printf(" persistent: %db\n", FIELD_SIZEOF(persistent)); + printf(" part1: %db\n", FIELD_SIZEOF(part1)); + printf(" part2: %db\n", FIELD_SIZEOF(part2)); + printf(" part3: %db\n", FIELD_SIZEOF(part3)); + printf(" part4: %db\n", FIELD_SIZEOF(part4)); + printf(" part5: %db\n", FIELD_SIZEOF(part5)); + printf(" part6: %db\n", FIELD_SIZEOF(part6)); + } +// =========================================== +#define FIELD(INSIDE_STRUCT, FIELDNAME) \ + CUTLASS_DEVICE auto& FIELDNAME() { \ + return INSIDE_STRUCT.FIELDNAME; \ + } + + FIELD(persistent, di) + FIELD(part1, lse_i) + FIELD(part1, mm_qk_k) + FIELD(part1, mm_qk_q) + // FIELD(part2, bias) + FIELD(part2, attn_shared_storage) + FIELD(part2, mm_gradV) + FIELD(part2, gradV_epilogue) + FIELD(part3, mm_doivj) + FIELD(part3, gradB_epilogue) + FIELD(part4, tmpT_shared_storage) + FIELD(part4, tmp_shared_storage) + FIELD(part4, mm_gradQ) + FIELD(part4, gradQ_epilogue) + FIELD(part4, gradQ_epilogue_lastIter) + FIELD(part5, mm_gradK) + FIELD(part5, gradK_epilogue) + FIELD(part6, gradK_epilogue_final) + FIELD(part6, gradV_epilogue_final) + }; + + using SharedStorage = typename cutlass::platform::conditional< + kPreload, + SharedStoragePrologue, + SharedStorageNoPrologue>::type; + + struct OutputFragments { + typename MatmulGradV::Mma::FragmentC gradV; + typename MatmulGradK::Mma::FragmentC gradK; + + CUTLASS_DEVICE void clear() { + gradV.clear(); + gradK.clear(); + } + }; + + static bool __host__ check_supported(Params const& p) { + CHECK_ALIGNED_PTR(p.query_ptr, kMinimumAlignment); + CHECK_ALIGNED_PTR(p.key_ptr, kMinimumAlignment); + CHECK_ALIGNED_PTR(p.value_ptr, kMinimumAlignment); + CHECK_ALIGNED_PTR(p.output_ptr, kMinimumAlignment); + CHECK_ALIGNED_PTR(p.grad_output_ptr, kMinimumAlignment); + + NATTEN_CHECK( + p.head_dim % kMinimumAlignment == 0, + "query is not correctly aligned (strideM)"); + NATTEN_CHECK( + p.head_dim % kMinimumAlignment == 0, + "key is not correctly aligned (strideM)"); + NATTEN_CHECK( + p.head_dim_value % kMinimumAlignment == 0, + "value is not correctly aligned (strideM)"); + NATTEN_CHECK(p.head_dim > 0, "Invalid value for `head_dim`"); + NATTEN_CHECK(p.head_dim_value > 0, "Invalid value for `head_dim_value`"); + // NATTEN_CHECK(p.num_queries > 0, "Invalid value for `num_queries`"); + // NATTEN_CHECK(p.num_keys > 0, "Invalid value for `num_keys`"); + NATTEN_CHECK(p.num_heads > 0, "Invalid value for `num_heads`"); + NATTEN_CHECK(p.num_batches > 0, "Invalid value for `num_batches`"); + NATTEN_CHECK(p.head_dim <= kMaxK, "kMaxK: Expected `head_dim < kMaxK`"); + NATTEN_CHECK( + p.head_dim_value <= kMaxK, "kMaxK: Expected `head_dim_value < kMaxK`"); + NATTEN_CHECK( + kEnableSplitKeys || p.num_splits_key.prod32() == 1, + "SplitKeys is disabled"); + // NATTEN_CHECK( + // p.num_splits_key > 0, "Invalid `num_splits_key` (expected >0)"); + + auto cta_shape_y = + ceil_div_dim(ceil_div_dim(p.num_queries, p.dilation), p.key_tile_shape); + NATTEN_CHECK( + // p.num_splits_key <= cutlass::ceil_div(p.num_keys, kBlockSizeJ), + // p.num_splits_key <= cta_shape_y.prod32(), + is_coord_less_than_or_equal_to(p.num_splits_key, cta_shape_y), + "Number of KV splits exceeds the number of tiles!"); + return true; + } + + static CUTLASS_DEVICE void attention_kernel(Params p) { + extern __shared__ char smem_buffer[]; + SharedStorage& shared_storage = *((SharedStorage*)smem_buffer); + + uint16_t thread_id = threadIdx.x; + uint8_t warp_id = gemm_kernel_utils::warp_uniform(thread_id / 32); + uint8_t lane_id = thread_id % 32; + + // int32_t key_start_int = p.split_key_device() * kBlockSizeJ; + // auto num_key_tiles = ceil_div_dim(p.num_queries_post_partitioning, + // p.key_tile_shape); auto key_start_idx = + // map_index_to_coord(p.split_key_device(), num_key_tiles); + auto key_start = p.split_key_device_dim() * p.key_tile_shape; + // if (key_start_int >= p.num_tokens) { + if (!is_coord_within_upper_bound( + key_start, p.num_queries_post_partitioning)) { + return; + } + + bool is_first_key = + p.num_splits_key_device() > 1 ? p.split_key_device() == 0 : true; + + // auto key_start = map_index_to_coord(key_start_int, + // p.num_queries_post_partitioning); + + if (kPrologueQK) { + auto query_start = getQueryStart(p, key_start); + prologueQkNextIteration( + shared_storage, p, query_start, key_start, warp_id, lane_id); + } + + // Computes (dO*out).sum(-1) and writes it to `p.delta_ptr` + if constexpr (kKernelComputesDelta) { + constexpr int kOptimalElements = + 128 / cutlass::sizeof_bits::value; + + // NOTE(alih): can't we parallelize this the same way we do KV? + // we'd just need a global fence? + + if (p.head_dim_value % kOptimalElements == 0) { + for (Dim query_start = Dim(); is_coord_within_upper_bound( + query_start, p.num_queries_post_partitioning); + increment_tile( + query_start, + p.query_tile_shape, + p.num_queries_post_partitioning)) { + computeDelta(p, query_start, warp_id, lane_id); + } + } else { + for (Dim query_start = Dim(); is_coord_within_upper_bound( + query_start, p.num_queries_post_partitioning); + increment_tile( + query_start, + p.query_tile_shape, + p.num_queries_post_partitioning)) { + computeDelta<1>(p, query_start, warp_id, lane_id); + } + } + __syncthreads(); + //} + } + + OutputFragments output_frags; + + CUTLASS_PRAGMA_UNROLL + for (; is_coord_within_upper_bound( + key_start, p.num_queries_post_partitioning); + increment_tile( + key_start, + p.key_tile_shape * p.num_splits_key, + /*lower_bound=*/p.split_key_device_dim() * p.key_tile_shape, + p.num_queries_post_partitioning)) { + output_frags.clear(); + + bool is_first_query = true; + bool is_same_key_tile = true; + auto next_key = key_start; + auto query_start = getQueryStart(p, key_start); + // while (next_key == key_start && query_start < p.num_queries) { + while (is_same_key_tile && + is_coord_within_upper_bound( + query_start, p.num_queries_post_partitioning)) { + // This line here + // vvvvvvvvvvvvvv + warp_id = gemm_kernel_utils::warp_uniform(warp_id); + // ^^^^^^^^^^^^^^ + // ... makes everything use less RF and be 10% faster. Why? + // I don't know. My theory is that it forces `nvcc` to + // re-compute indices, offsets etc... and not keep them + // from the previous iteration, which prevents MASSIVE + // register spilling. + + processBlockIJ( + shared_storage, + output_frags, + p, + query_start, + key_start, + is_first_query, + is_first_key, + warp_id, + lane_id); + + is_first_query = false; + is_first_key = false; + Dim next_query; + is_same_key_tile = + incrIteration(p, query_start, key_start, next_query, next_key); + query_start = next_query; + } + if (kOutputInRF) { + writeFragsToGmem( + shared_storage, output_frags, p, key_start, warp_id, lane_id); + //} else if (getQueryStart(p, key_start) >= p.num_queries) { + // zfillGradKV( + // p, key_start, warp_id, lane_id); + } + __syncthreads(); + } + } + + // This never happens in NA :D . + // static CUTLASS_DEVICE void zfillGradKV( + // Params const& p, + // int32_t key_start, + // uint8_t warp_id, + // uint8_t lane_id) { + // constexpr int kThreadsPerKey = 8; + // constexpr int kParallelKeys = kNumThreads / kThreadsPerKey; + // static_assert(kBlockSizeJ % kParallelKeys == 0, ""); + // // This function is not really optimized, but should rarely be used + // // It's only used when some keys are "useless" and don't attend to + // // any query, due to causal masking + + // int thread_id = 32 * warp_id + lane_id; + // int k_shift = lane_id % kThreadsPerKey; + + // CUTLASS_PRAGMA_UNROLL + // for (int j = 0; j < kBlockSizeJ; j += kParallelKeys) { + // int key = key_start + j + (thread_id / kThreadsPerKey); + // if (key >= p.num_keys) { + // continue; + // } + // auto gv_ptr = p.grad_value_ptr + key * p.gV_strideM(); + // auto gk_ptr = p.grad_key_ptr + key * p.gK_strideM(); + + // for (int k = k_shift; k < p.head_dim_value; k += kThreadsPerKey) { + // gv_ptr[k] = scalar_t(0); + // } + // for (int k = k_shift; k < p.head_dim; k += kThreadsPerKey) { + // gk_ptr[k] = scalar_t(0); + // } + // } + //} + + static CUTLASS_DEVICE void processBlockIJ( + SharedStorage& shared_storage, + OutputFragments& output_frags, + Params& p, + Dim query_start, + Dim key_start, + bool isFirstQuery, + bool isFirstKey, + uint8_t warp_id, + uint8_t lane_id) { + cutlass::MatrixCoord no_offset{0, 0}; + accum_t scale = p.scale; + int16_t thread_id = 32 * warp_id + lane_id; + + // auto key_start_int = map_coord_to_index(key_start, + // p.num_queries_post_partitioning); + + auto rematerializeThreadIds = [&]() { + // Prevents `nvcc` from keeping values deduced from + // `thread_id`, `warp_id`, ... in RF - to reduce register pressure + warp_id = gemm_kernel_utils::warp_uniform(thread_id / 32); + lane_id = thread_id % 32; + thread_id = 32 * warp_id + lane_id; + }; + + // bool isFirstQuery = (query_start == getQueryStart(p, key_start)); + Dim next_query, next_key; + bool isLastQuery = + !incrIteration(p, query_start, key_start, next_query, next_key); + // bool isLastQuery = next_key != key_start; + // PRINT_T0("[b:(%d, %d, %d)] q=%d k=%d FIRST=%d LAST=%d next: (%d, %d)", + // int(blockIdx.x), int(blockIdx.y), int(blockIdx.z), + // query_start, key_start, int(isFirstQuery), int(isLastQuery), + // int(next_query), int(next_key) + // ); + + // int32_t num_queries_in_block = + // gemm_kernel_utils::warp_uniform(cutlass::fast_min((int32_t)MatmulQK::Mma::Shape::kN, + // p.num_queries - query_start)); int32_t num_keys_in_block = + // gemm_kernel_utils::warp_uniform(cutlass::fast_min((int32_t)MatmulQK::Mma::Shape::kM, + // p.num_keys - key_start)); + // TODO(alih): make warp uniform? + // auto num_queries_in_block = + // cutlass::fast_min((int32_t)MatmulQK::Mma::Shape::kN, p.num_queries - + // query_start); + auto num_queries_in_block = fast_min( + p.query_tile_shape, p.num_queries_post_partitioning - query_start); + // auto num_keys_in_block = + // cutlass::fast_min((int32_t)MatmulQK::Mma::Shape::kM, p.num_keys - + // key_start); + auto num_keys_in_block = + fast_min(p.key_tile_shape, p.num_queries_post_partitioning - key_start); + + auto num_queries_in_block_int = num_queries_in_block.prod32(); + auto num_keys_in_block_int = num_keys_in_block.prod32(); + + // NOTE(alih): we fetch lse_i here as well as d_i. + // We do this to avoid having to use the shmem epilogue iterator + // to read lse, since lse is not aligned and needs to be gathered + // for a GETT, and given that there's a single scalar per row element, + // having an iterator is probably unnecessarily complicated. + accum_t lsei_rf = accum_t(0); + accum_t di_rf = accum_t(0); + if (thread_id < kBlockSizeI) { + auto query_idx = query_start + + map_index_to_coord((int32_t)thread_id, num_queries_in_block); + auto query_offset = (query_idx * p.lse_strideM).sum(); + if (is_coord_within_upper_bound( + query_idx, p.num_queries_post_partitioning)) { + di_rf = p.delta_ptr[query_offset]; + lsei_rf = p.logsumexp_ptr[query_offset]; + } + shared_storage.lse_i()[thread_id] = lsei_rf; + shared_storage.di()[thread_id] = di_rf; + } + // NOTE(alih): we don't need to syncthreads here; + // mma will do that before anything tries to access + // lse_i, and di is not used until after part1. + //__syncthreads(); + // end NOTE + // end NOTE + + auto prologueGradV = [&](int col) { + typename MatmulGradV::Mma::IteratorB iterator_dO( + //{int32_t(p.gO_strideM)}, + {p.o_strideM, num_queries_in_block}, + // p.grad_output_ptr + query_start * p.gO_strideM + col, + p.grad_output_ptr + (query_start * p.o_strideM).sum() + col, + //{num_queries_in_block, p.head_dim_value - col}, + num_queries_in_block, + p.head_dim_value - col, + thread_id, + no_offset); + MatmulGradV::Mma::prologue( + shared_storage.mm_gradV(), + iterator_dO, + thread_id, + num_queries_in_block_int); + }; + auto prologueGradQ = [&](int col) { + typename MatmulGradQ::Mma::IteratorB iterator_K( + //{int32_t(p.k_strideM)}, + {p.k_strideM, num_keys_in_block}, + // p.key_ptr + key_start * p.k_strideM + col, + p.key_ptr + (key_start * p.k_strideM).sum() + col, + //{num_keys_in_block, p.head_dim - col}, + num_keys_in_block, + p.head_dim - col, + thread_id, + no_offset); + MatmulGradQ::Mma::prologue( + shared_storage.mm_gradQ(), + iterator_K, + thread_id, + num_keys_in_block_int); + }; + auto prologueGradK = [&](int col) { + typename MatmulGradK::Mma::IteratorB iterator_Q( + //{int32_t(p.q_strideM)}, + {p.q_strideM, num_queries_in_block}, + // p.query_ptr + query_start * p.q_strideM + col, + p.query_ptr + (query_start * p.q_strideM).sum() + col, + //{num_queries_in_block, p.head_dim - col}, + num_queries_in_block, + p.head_dim - col, + thread_id, + no_offset); + MatmulGradK::Mma::prologue( + shared_storage.mm_gradK(), + iterator_Q, + thread_id, + num_queries_in_block_int); + }; + auto prologueDOV = [&]() { + typename MatmulDOIVJ::Mma::IteratorA iterator_A( + //{int32_t(p.gO_strideM)}, + {p.o_strideM, num_queries_in_block}, + // p.grad_output_ptr + query_start * p.gO_strideM, + p.grad_output_ptr + (query_start * p.o_strideM).sum(), + num_queries_in_block, + p.head_dim_value, + // {num_queries_in_block, p.head_dim_value}, + thread_id, + no_offset); + typename MatmulDOIVJ::Mma::IteratorB iterator_B( + //{int32_t(p.v_strideM)}, + {p.v_strideM, num_keys_in_block}, + // p.value_ptr + key_start * p.v_strideM, + p.value_ptr + (key_start * p.v_strideM).sum(), + //{p.head_dim_value, num_keys_in_block}, + num_keys_in_block, + p.head_dim_value, + thread_id, + no_offset); + MatmulDOIVJ::Mma::prologue( + shared_storage.mm_doivj(), + iterator_A, + iterator_B, + thread_id, + p.head_dim_value); + }; + + ///////////////////////////////////////////////////////////////////////////////////////////////// + // MatmulQK + ///////////////////////////////////////////////////////////////////////////////////////////////// + { + using Mma = typename MatmulQK::Mma; + + // cutlass::gemm::GemmCoord problem_size( + // num_keys_in_block_int, + // num_queries_in_block_int, + // p.head_dim // k + //); + + // k_j + typename Mma::IteratorA iterator_A( + //{int32_t(p.k_strideM)}, + {p.k_strideM, num_keys_in_block}, + // p.key_ptr + key_start * p.k_strideM, + p.key_ptr + (key_start * p.k_strideM).sum(), + // {problem_size.m(), problem_size.k()}, + num_keys_in_block, + p.head_dim, + thread_id, + no_offset); + + // q_i.transpose(-2, -1) + typename Mma::IteratorB iterator_B( + //{int32_t(p.q_strideM)}, + {p.q_strideM, num_queries_in_block}, + p.query_ptr + (query_start * p.q_strideM).sum(), + //{problem_size.k(), problem_size.n()}, + num_queries_in_block, + // problem_size.k(), + p.head_dim, + thread_id, + no_offset); + + Mma mma( + shared_storage.mm_qk_k(), + shared_storage.mm_qk_q(), + thread_id, + warp_id, + lane_id); + + typename Mma::FragmentC accum; + + accum.clear(); + + // auto gemm_k_iterations = + // (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK; + auto gemm_k_iterations = + (p.head_dim + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Compute threadblock-scoped matrix multiply-add + mma.set_prologue_done(kPrologueQK); + mma.set_zero_outside_bounds(/*!skipBoundsChecks*/ true); + mma(gemm_k_iterations, accum, iterator_A, iterator_B, accum); + accum = cutlass::multiplies()(scale, accum); + + // Epilogue: add LSE + exp and store that to our shared memory buffer + // shmem <- (matmul_result - + // logsumexp[i_start:i_end].unsqueeze(1)).exp() + int warp_idx_mn_0 = + warp_id % (Mma::Base::WarpCount::kM * Mma::Base::WarpCount::kN); + auto output_tile_coords = cutlass::MatrixCoord{ + warp_idx_mn_0 % Mma::Base::WarpCount::kM, + warp_idx_mn_0 / Mma::Base::WarpCount::kM}; + + auto lane_offset = MatmulQK::AccumLambdaIterator::get_lane_offset( + lane_id, warp_id, output_tile_coords); + + if (not p.is_fully_block_sparse) { + // Neighborhood Attention masking + Dim first_col, query_bound, row_idx; + + auto na_mask = + NAMask(p.kernel_size, p.stride, p.num_queries_post_partitioning); + + MatmulQK::AccumLambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) { + row_idx = + map_index_to_coord((int32_t)accum_m, num_keys_in_block) + + key_start; + first_col = na_mask.get_backward_window_start(row_idx); + query_bound = + na_mask.get_backward_window_end(row_idx) - first_col; + }, + [&](int accum_m, int accum_n, int idx) { + auto col = + map_index_to_coord((int32_t)accum_n, num_queries_in_block) + + query_start; + // NOTE(alih): Checking whether the current q coordinate is out of + // bounds is something I had to do for SM70 and SM50 kernels. + // Without it, a lot of query tile configurations but the default + // failed. I guess their GEMMs align the number of iters up to the + // thread-level K shape. While we can avoid this condition on SM80 + // kernels, I don't know how safe that is, so I'll just keep it + // for all of them. + if (accum_m >= num_keys_in_block_int || +#if __CUDA_ARCH__ < 800 + accum_n >= num_queries_in_block_int || +#endif + !is_coord_within_bounds_nn(col - first_col, query_bound)) { + accum[idx] = + -cutlass::platform::numeric_limits::infinity(); + } else { + accum[idx] -= shared_storage.lse_i()[accum_n]; + } + }, + [&](int accum_m) {}); + + } else if (p.has_q_padding) { + // Residual masking + + MatmulQK::AccumLambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) {}, + [&](int accum_m, int accum_n, int idx) { + auto query_coord = + map_index_to_coord((int32_t)accum_n, num_queries_in_block) + + query_start; + +#if __CUDA_ARCH__ < 800 + if (accum_m >= num_keys_in_block_int || + accum_n >= num_queries_in_block_int || + !is_coord_less_than_or_equal_to( + query_coord, p.num_queries_post_partitioning)) { +#else + if (!is_coord_less_than_or_equal_to( + query_coord, p.num_queries_post_partitioning)) { +#endif + accum[idx] = + -cutlass::platform::numeric_limits::infinity(); + } else { + accum[idx] -= shared_storage.lse_i()[accum_n]; + } + }, + [&](int accum_m) {}); + } else { + // No masking; just subtract LSE + + MatmulQK::AccumLambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) {}, + [&](int accum_m, int accum_n, int idx) { +#if __CUDA_ARCH__ < 800 + if (accum_m >= num_keys_in_block_int || + accum_n >= num_queries_in_block_int) { + accum[idx] = + -cutlass::platform::numeric_limits::infinity(); + } else { +#endif + accum[idx] -= shared_storage.lse_i()[accum_n]; +#if __CUDA_ARCH__ < 800 + } +#endif + }, + [&](int accum_m) {}); + } + + __syncthreads(); + if (kPrologueGV) { + prologueGradV(0); + } + if (kPrologueDOV) { + prologueDOV(); + } + + // MatmulQK::B2bGemm::accumApplyLSEToSmem( + MatmulQK::B2bGemm::accumApplyExpToSmem( + shared_storage.attn_shared_storage(), + accum, + // p.logsumexp_ptr + (query_start * p.lse_strideM).sum(), + // p.lse_strideM, //<- stride + // num_queries_in_block, //<- extent + ////problem_size.n(), + thread_id, + warp_id, + lane_id, + output_tile_coords); + + __syncthreads(); + } + rematerializeThreadIds(); + + ///////////////////////////////////////////////////////////////////////////////////////////////// + // GradV matmul + // + // grad_v[j_start:j_end] += attn_T @ do_i + ///////////////////////////////////////////////////////////////////////////////////////////////// + constexpr bool kSingleIterationGradV = + kMaxK <= MatmulGradV::ThreadblockShape::kN; + for (int col = 0; col < (kSingleIterationGradV ? 1 : p.head_dim_value); + col += MatmulGradV::ThreadblockShape::kN) { + using Mma = typename MatmulGradV::Mma; + using AccumTileGmem = typename MatmulGradV::AccumTileGmem; + + // cutlass::gemm::GemmCoord problem_size( + // num_keys_in_block_int, + // p.head_dim_value - col, + // num_queries_in_block_int); + auto createEpilogueIter = [&]() { + return typename MatmulGradV::OutputTileIterator( + // typename MatmulGradV::OutputTileIterator::Params{p.gV_strideM()}, + // p.grad_value_ptr + key_start * p.gV_strideM() + col, + //{num_keys_in_block, p.head_dim_value - col}, + typename MatmulGradV::OutputTileIterator::Params{ + p.v_strideM, num_keys_in_block}, + p.grad_value_ptr + (key_start * p.v_strideM).sum() + col, + num_keys_in_block, + p.head_dim_value - col, + thread_id); + }; + typename Mma::IteratorB iterator_B( + //{int32_t(p.gO_strideM)}, + // p.grad_output_ptr + query_start * p.gO_strideM + col, + //{num_queries_in_block, p.head_dim_value - col}, + {p.o_strideM, num_queries_in_block}, + p.grad_output_ptr + (query_start * p.o_strideM).sum() + col, + num_queries_in_block, + p.head_dim_value - col, + thread_id, + no_offset); + + // dVj += Pij.T @ dOi + Mma mma( + // operand A: Pij.T + shared_storage.attn_shared_storage().accum_ref(), + // operand B: dOi - which was loaded into shared memory previously + // when we computed dVj + shared_storage.mm_gradV().operand_B_ref(), + thread_id, + warp_id, + lane_id); + + int storage_id = col / MatmulGradV::ThreadblockShape::kN; + AccumTileGmem gmem_tile{ + p.workspace_gv + storage_id * AccumTileGmem::kElementsStored}; + if (!kOutputInRF) { + if (isFirstQuery || !kNeedsAccumGradV) { + output_frags.gradV.clear(); + } else { + gmem_tile.load(output_frags.gradV, thread_id); + } + } + mma.set_prologue_done(kPrologueGV); + + // auto gemm_k_iterations = + // (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK; + auto gemm_k_iterations = + (num_queries_in_block_int + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Compute threadblock-scoped matrix multiply-add + __syncthreads(); + + mma(gemm_k_iterations, + output_frags.gradV, + iterator_B, + output_frags.gradV); + __syncthreads(); + if (kPrologueGV && !kSingleIterationGradV && + col + MatmulGradV::ThreadblockShape::kN < p.head_dim_value) { + prologueGradV(col + MatmulGradV::ThreadblockShape::kN); + } + + if (!kOutputInRF) { + if (kNeedsAccumGradV && !isLastQuery) { + gmem_tile.store(output_frags.gradV, thread_id); + } else { + accumulateInGmem( + shared_storage.gradV_epilogue(), + output_frags.gradV, + createEpilogueIter(), + isFirstQuery || kNeedsAccumGradV, + warp_id, + lane_id); + } + } + } + + //// NOTE(alih): we fetch d_i here instead of at the beginning; + //// because we use the same storage in shared memory for lse_i; + //// We do this to avoid having to use the shmem epilogue iterator + //// to read lse, since lse is not aligned and needs to be gathered + //// for NA. + // accum_t di_rf = accum_t(0); + // if (thread_id < kBlockSizeI) { + // auto query_idx = query_start + map_index_to_coord((int32_t)thread_id, + // num_queries_in_block); auto query_offset = (query_idx * + // p.lse_strideM).sum(); if (is_coord_within_upper_bound(query_idx, + // p.num_queries_post_partitioning)) { + // //if (query_start + thread_id < p.num_queries) { + // // di_rf = p.delta_ptr[query_start + thread_id]; + // di_rf = p.delta_ptr[query_offset]; + // } + // shared_storage.di()[thread_id] = di_rf; + // } + //// end NOTE + + __syncthreads(); + + ///////////////////////////////////////////////////////////////////////////////////////////////// + // MatmulDOIVJ + ///////////////////////////////////////////////////////////////////////////////////////////////// + { + using Mma = typename MatmulDOIVJ::Mma; + // do_i + typename Mma::IteratorA iterator_A( + //{int32_t(p.gO_strideM)}, + // p.grad_output_ptr + query_start * p.gO_strideM, + //{num_queries_in_block, p.head_dim_value}, + {p.o_strideM, num_queries_in_block}, + p.grad_output_ptr + (query_start * p.o_strideM).sum(), + num_queries_in_block, + p.head_dim_value, + thread_id, + no_offset); + + // v_j.transpose(-2, -1) + typename Mma::IteratorB iterator_B( + //{int32_t(p.v_strideM)}, + // p.value_ptr + key_start * p.v_strideM, + //{p.head_dim_value, num_keys_in_block}, + {p.v_strideM, num_keys_in_block}, + p.value_ptr + (key_start * p.v_strideM).sum(), + num_keys_in_block, + p.head_dim_value, + thread_id, + no_offset); + + Mma mma(shared_storage.mm_doivj(), thread_id, warp_id, lane_id); + mma.set_prologue_done(kPrologueDOV); + mma.set_zero_outside_bounds(/*!skipBoundsChecks*/ true); + + typename Mma::FragmentC accum; + + accum.clear(); + + auto gemm_k_iterations = + (p.head_dim_value + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Compute threadblock-scoped matrix multiply-add + mma(gemm_k_iterations, accum, iterator_A, iterator_B, accum); + __syncthreads(); + if (kPrologueGQ) { + prologueGradQ(0); + } + if (kPrologueGK) { + prologueGradK(0); + } + + int warp_idx_mn_0 = + warp_id % (Mma::Base::WarpCount::kM * Mma::Base::WarpCount::kN); + auto output_tile_coords = cutlass::MatrixCoord{ + warp_idx_mn_0 % Mma::Base::WarpCount::kM, + warp_idx_mn_0 / Mma::Base::WarpCount::kM}; + // TODO: This must be terribly inefficient. There must be a better way + // tmp [RF] <- (accum [RF] - Di [smem] ) * attn_T.T [smem] + // attn_shared_storage [smem] <- tmp.T + // tmp_shared_storage [smem] <- tmp + { + using LambdaIterator = typename MatmulDOIVJ::AccumLambdaIterator; + auto lane_offset = LambdaIterator::get_lane_offset( + lane_id, warp_id, output_tile_coords); + + auto attn_T = shared_storage.attn_shared_storage().accum_ref(); + + accum_t current_di; + // dSij = (dPij - Di) * Pij + LambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) { current_di = shared_storage.di()[accum_m]; }, + [&](int accum_m, int accum_n, int idx) { + // TODO: Otherwise we can get nans as we + // might have infs here (only seen on f16 tho) + if (accum_m < num_queries_in_block_int && + accum_n < num_keys_in_block_int) { + accum_t attn = attn_T.at({accum_n, accum_m}); + accum[idx] = (accum[idx] - current_di) * attn; + } else { + accum[idx] = 0; + } + }, + [&](int accum_m) { + + }); + + // store bias gradient tile dBij to global memory, + // where dBij = dSij = Pij * (dPij - Di) + // if (p.grad_bias_ptr != nullptr) { + // typename MatmulDOIVJ::BiasGradEpilogue::OutputTileIterator + // output_iter( + // typename + // MatmulDOIVJ::BiasGradEpilogue::OutputTileIterator:: + // Params{p.gB_strideM}, + // // grad_bias_ptr is offset to point at beginning of + // // matrix of shape (queries, keys) for a given + // // (batch_id, head_id) the pointer arithmetic here produces + // // a pointer to the start of the current tile within that + // // matrix + // p.grad_bias_ptr + query_start * p.gB_strideM + key_start, + // {num_queries_in_block, num_keys_in_block}, + // thread_id); + + // // no-op epilogue operator - just casting and storing contents of + // // accum to global memory + // typename MatmulDOIVJ::BiasGradEpilogue::OutputOp output_op({1, 1}); + // typename MatmulDOIVJ::BiasGradEpilogue epilogue( + // shared_storage.gradB_epilogue(), thread_id, warp_id, lane_id); + // epilogue(output_op, output_iter, accum, output_iter); + // } + + accum = accum * scale; + +#if 0 + PRINT_B0_T0("(doivj - di) * attn * scale"); + print_warp_accum(accum, lane_offset, 4, 4); +#endif + + __syncthreads(); + if (!MatmulGradK::DefaultMmaFromSmem::kIsTransposedA) { + auto tmpT = shared_storage.tmpT_shared_storage().accum_ref(); + // attn <- attn_T.T + LambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) {}, + [&](int accum_m, int accum_n, int idx) { + tmpT.at({accum_n, accum_m}) = scalar_t(accum[idx]); + }, + [&](int accum_m) {}); + } + } + + MatmulDOIVJ::B2bGemm::accumToSmem( + shared_storage.tmp_shared_storage(), + accum, + lane_id, + output_tile_coords); + __syncthreads(); + } + // Force `nvcc` to recompute values that depend on the variables just below + // to use less RF and prevent some spilling + p.head_dim = gemm_kernel_utils::warp_uniform(p.head_dim); + // p.k_strideM = gemm_kernel_utils::warp_uniform(p.k_strideM); + rematerializeThreadIds(); + + ///////////////////////////////////////////////////////////////////////////////////////////////// + // GradQ matmul + // + // grad_q[i_start:i_end] += tmp @ k_j + ///////////////////////////////////////////////////////////////////////////////////////////////// + // Skip the loop & associated branches if we know at compile time the number + // of iterations + constexpr bool kSingleIterationGradQ = + kMaxK <= MatmulGradQ::ThreadblockShape::kN; + for (int col = 0; col < (kSingleIterationGradQ ? 1 : p.head_dim); + col += MatmulGradQ::ThreadblockShape::kN) { + using Mma = typename MatmulGradQ::Mma; + using AccumTileGmem = typename MatmulGradQ::AccumTileGmem; + + // cutlass::gemm::GemmCoord problem_size( + // num_queries_in_block_int, + // //false ? MatmulGradQ::ThreadblockShape::kN : p.head_dim - col, + // p.head_dim - col, + // num_keys_in_block_int); + + // k_j + typename Mma::IteratorB iterator_B( + //{int32_t(p.k_strideM)}, + // p.key_ptr + key_start * p.k_strideM + col, + //{problem_size.k(), problem_size.n()}, + {p.k_strideM, num_keys_in_block}, + p.key_ptr + (key_start * p.k_strideM).sum() + col, + num_keys_in_block, + p.head_dim - col, + thread_id, + no_offset); + + auto a = shared_storage.tmp_shared_storage().accum_ref(); + Mma mma( + // operand A: dSij + shared_storage.tmp_shared_storage().accum_ref(), + // operand B: Kj + shared_storage.mm_gradQ().operand_B_ref(), + thread_id, + warp_id, + lane_id); + + typename Mma::FragmentC accum; + + int col_id = col / MatmulGradQ::ThreadblockShape::kN; + int num_cols = kSingleIterationGradQ + ? 1 + : gemm_kernel_utils::ceil_div( + p.head_dim, MatmulGradQ::ThreadblockShape::kN); + // auto query_block_int = map_coord_to_index(div_dim(query_start, + // p.query_tile_shape), p.cta_shape_x); + auto query_block_int = map_coord_to_index( + div_dim(query_start, p.query_tile_shape), + ceil_div_dim(p.num_queries_post_partitioning, p.query_tile_shape)); + int storage_id = (col_id + query_block_int * num_cols); + + if (p.num_splits_key_device() > 1) { + AtomicLock::acquire( + &p.workspace_gq[storage_id].lock, + p.split_key_device() + 1, + thread_id); + // Make sure we can see other block's output + __threadfence(); + } + + AccumTileGmem gmem_tile{&p.workspace_gq[storage_id].buffer[0]}; + if (!kNeedsAccumGradQ || + //(p.num_splits_key_device() == 1 && key_start_int == 0)) { + (p.num_splits_key_device() == 1 && isFirstKey)) { + // if we know we are the first to access it, we know it's only zeros. + // Avoids a load from gmem (and gmem init as well) + accum.clear(); + } else { + gmem_tile.load(accum, thread_id); + } + + // auto gemm_k_iterations = + // (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK; + auto gemm_k_iterations = + (num_keys_in_block_int + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Compute threadblock-scoped matrix multiply-add + __syncthreads(); + mma.set_prologue_done(kPrologueGQ); + mma(gemm_k_iterations, accum, iterator_B, accum); + __syncthreads(); + bool isLastColumn = kSingleIterationGradQ || + (col + MatmulGradQ::ThreadblockShape::kN >= p.head_dim); + if (kPrologueGQ && !isLastColumn) { + prologueGradQ(col + MatmulGradQ::ThreadblockShape::kN); + } + + bool isLast = [&]() { + // int32_t next_key = key_start + p.num_splits_key_device() * + // kBlockSizeJ; + // NOTE(alih): we purposefully don't increment tile here! + // The reason is we need to make sure this CTA won't ever revisit + // this query tile in either direction (2D and 3D). + auto next_key = key_start + p.key_tile_shape * p.num_splits_key; + + auto na_mask = + NAMask(p.kernel_size, p.stride, p.num_queries_post_partitioning); + auto last_query_for_block = fast_min( + query_start + p.query_tile_shape, + p.num_queries_post_partitioning) - + 1; + + auto last_possible_key_tile = + align_dim_up(p.num_queries_post_partitioning, p.key_tile_shape) - + p.key_tile_shape; + // auto last_key_block_for_query_block = + // fast_min(align_dim_down(na_mask.get_window_end_( + // last_query_for_block) - 1, p.key_tile_shape), + // last_possible_key_tile); + auto last_key_block_for_query_block = + na_mask.get_last_key_block_for_query_block( + last_query_for_block, p.key_tile_shape, last_possible_key_tile); + + // auto next_key = key_start + p.key_tile_shape * + // p.num_splits_key_device(); if (p.num_keys <= next_key) { + if (is_coord_within_upper_bound( + last_key_block_for_query_block, next_key)) { + return true; + } + if (is_coord_greater_than_or_equal_to( + key_start, last_key_block_for_query_block)) { + return true; + } + return false; + }(); + // Output results + if (p.num_splits_key_device() > 1) { + int32_t numAddsSoFar = -1; + if (isLast && thread_id == 0) { + numAddsSoFar = atomicAdd(&p.workspace_gq[storage_id].counter, 1) + + 1; // `atomicAdd` returns the old value + } + auto numParallelBlocksForQuery = + getNumParallelBlocksForQuery(p, query_start); + isLast = __syncthreads_or(numAddsSoFar == numParallelBlocksForQuery); +#ifdef NATTEN_ENABLE_DEVICE_SIDE_ASSERTIONS + assert(numAddsSoFar <= numParallelBlocksForQuery); +#endif + } + if (kNeedsAccumGradQ && !isLast) { + gmem_tile.store(accum, thread_id); + if (p.num_splits_key_device() > 1) { + // Make sure everyone wrote before we release the lock + __threadfence(); + __syncthreads(); + AtomicLock::release(&p.workspace_gq[storage_id].lock, thread_id); + } + } else { + // NOTE: We're not releasing the lock because no one is expected + // to come after us (we're the last one to write) + typename MatmulGradQ::OutputTileIterator output_it( + // typename MatmulGradQ::OutputTileIterator::Params{p.gQ_strideM()}, + // p.grad_query_ptr + query_start * p.gQ_strideM() + col, + //{problem_size.m(), problem_size.n()}, + typename MatmulGradQ::OutputTileIterator::Params{ + p.q_strideM, num_queries_in_block}, + p.grad_query_ptr + (query_start * p.q_strideM).sum() + col, + num_queries_in_block, + p.head_dim - col, + thread_id); + // if `direct_store` is True, we store to gmem (`*gmem = accum`) + // otherwise, we accumulate in gmem (`*gmem = *gmem + accum`) + // If we know ahead of time when we will write for the first time + // we can: + // (1) Avoid an additional memory read + // (2) Avoid the cost of initializing memory to 0 + // bool direct_store = kNeedsAccumGradQ || key_start_int == 0 || + bool direct_store = + kNeedsAccumGradQ || isFirstKey || (p.num_splits_key_device() > 1); + accumulateInGmem( + isLastColumn ? shared_storage.gradQ_epilogue_lastIter() + : shared_storage.gradQ_epilogue(), + accum, + output_it, + direct_store, + warp_id, + lane_id); + } + } + ///////////////////////////////////////////////////////////////////////////////////////////////// + // GradK matmul + // + // grad_k[i_start:i_end] += tmp.transpose(-2, -1) @ q_i + ///////////////////////////////////////////////////////////////////////////////////////////////// + rematerializeThreadIds(); + + constexpr bool kSingleIterationGradK = + kMaxK <= MatmulGradK::ThreadblockShape::kN; + for (int col = 0; col < (kSingleIterationGradK ? 1 : p.head_dim); + col += MatmulGradK::ThreadblockShape::kN) { + using Mma = typename MatmulGradK::Mma; + using AccumTileGmem = typename MatmulGradQ::AccumTileGmem; + + // cutlass::gemm::GemmCoord problem_size( + // num_keys_in_block_int, + // //false ? MatmulGradK::ThreadblockShape::kN : p.head_dim - col, + // p.head_dim - col, + // num_queries_in_block_int); + auto createEpilogueIter = [&]() { + return typename MatmulGradK::OutputTileIterator( + // typename MatmulGradK::OutputTileIterator::Params{p.gK_strideM()}, + // p.grad_key_ptr + key_start * p.gK_strideM() + col, + //{num_keys_in_block_int, + // false ? MatmulGradK::ThreadblockShape::kN : p.head_dim - col}, + typename MatmulGradK::OutputTileIterator::Params{ + p.k_strideM, num_keys_in_block}, + p.grad_key_ptr + (key_start * p.k_strideM).sum() + col, + num_keys_in_block, + p.head_dim - col, + thread_id); + }; + + // q_i + typename Mma::IteratorB iterator_B( + //{int32_t(p.q_strideM)}, + // p.query_ptr + query_start * p.q_strideM + col, + //{problem_size.k(), problem_size.n()}, + {p.q_strideM, num_queries_in_block}, + p.query_ptr + (query_start * p.q_strideM).sum() + col, + num_queries_in_block, + p.head_dim - col, + thread_id, + no_offset); + + auto getTmp = [&](int) { return &shared_storage.tmp_shared_storage(); }; + auto getTmpT = [&](int) { return &shared_storage.tmpT_shared_storage(); }; + // this is basically: + // opA = kIsTransposedA ? getTmp() : getTmpT(); + bool constexpr kIsTransposedA = + MatmulGradK::DefaultMmaFromSmem::kIsTransposedA; + auto& opA = *gemm_kernel_utils::call_conditional< + kIsTransposedA, + decltype(getTmp), + decltype(getTmpT)>::apply(getTmp, getTmpT, 0); + Mma mma( + // operand A: dSij.T + opA.accum_ref(), + // operand B: Qi + shared_storage.mm_gradK().operand_B_ref(), + thread_id, + warp_id, + lane_id); + + int storage_id = col / MatmulGradK::ThreadblockShape::kN; + AccumTileGmem gmem_tile{ + p.workspace + storage_id * AccumTileGmem::kElementsStored}; + if (!kOutputInRF) { + if (isFirstQuery || !kNeedsAccumGradK) { + output_frags.gradK.clear(); + } else { + gmem_tile.load(output_frags.gradK, thread_id); + } + } + mma.set_prologue_done(kPrologueGK); + + // auto gemm_k_iterations = + // (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK; + auto gemm_k_iterations = + (num_queries_in_block_int + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Compute threadblock-scoped matrix multiply-add + __syncthreads(); + + mma(gemm_k_iterations, + output_frags.gradK, + iterator_B, + output_frags.gradK); + __syncthreads(); + bool isLastColumn = kSingleIterationGradK || + col + MatmulGradK::ThreadblockShape::kN >= p.head_dim; + if (kPrologueGK && !isLastColumn) { + prologueGradK(col + MatmulGradK::ThreadblockShape::kN); + } + + if (kPrologueQK && isLastColumn) { + Dim next_query, next_key; + bool isLastQuery = + !incrIteration(p, query_start, key_start, next_query, next_key); + DISPATCH_BOOL( + // next_key != key_start, kForceReloadK, ([&]() { + isLastQuery, + kForceReloadK, + ([&]() { + prologueQkNextIteration( + shared_storage, p, next_query, next_key, warp_id, lane_id); + })); + } + + // Output results + if (!kOutputInRF) { + if (kNeedsAccumGradK && !isLastQuery) { + gmem_tile.store(output_frags.gradK, thread_id); + } else { + accumulateInGmem( + isLastColumn ? shared_storage.gradK_epilogue_final() + : shared_storage.gradK_epilogue(), + output_frags.gradK, + createEpilogueIter(), + isFirstQuery || kNeedsAccumGradK, + warp_id, + lane_id); + __syncthreads(); + } + } + } + } + + // NOTE (alih): Shifting query tiles was originally just intended for + // when there was no attention masking, which includes sliding windows. + // So we'll just comment out the whole logic relating to that. + // I couldn't come up with a way to extend this logic to sliding windows + // in time. + // TODO: can we shift query tiles as long as we know there's more than one + // parallel key tile? It should be possible because in our case we start with + // the assumption that we have the same number of queries and keys in gmem, + // but what's different is the number of query and key tiles necessary? + // static CUTLASS_HOST_DEVICE int32_t getQueryStartShift(Params const& p) { + // if (p.custom_mask_type == NoCustomMask && p.num_splits_key_device() > 1) { + // return (p.split_key_device() * kBlockSizeI) % getQueryEnd(p); + // } + // return 0; + //} + + // Iteration order logic + static CUTLASS_HOST_DEVICE Dim getQueryStart(Params const& p, Dim key_start) { + // return getSmallestQueryForKey(p, key_start) + getQueryStartShift(p); + return getSmallestQueryForKey(p, key_start); + } + + static CUTLASS_HOST_DEVICE Dim getQueryEnd(Params const& p) { + // return gemm_kernel_utils::align_up(p.num_queries, kBlockSizeI); + return align_dim_up(p.num_queries_post_partitioning, p.query_tile_shape); + } + + static CUTLASS_HOST_DEVICE Dim + getSmallestQueryForKey(Params const& p, Dim key_start) { + // if (p.custom_mask_type == NoCustomMask) { + // return 0; + // } + // int32_t shift = p.custom_mask_type == CausalFromBottomRight + // ? p.num_keys - p.num_queries + // : 0; + // int32_t window_size = + // p.window_size == 0 ? p.num_queries + p.num_keys : p.window_size; + + // auto last_key_for_block = + // cutlass::fast_min(key_start + kBlockSizeJ, p.num_keys) - 1; + // int first_query = key_start - shift; + // int last_query = last_key_for_block - shift + window_size - 1; + // auto last_key_for_block = fast_min( + // key_start + p.key_tile_shape, + // p.num_queries_post_partitioning) - 1; + auto na_mask = + NAMask(p.kernel_size, p.stride, p.num_queries_post_partitioning); + + auto first_query = na_mask.get_backward_window_start(key_start); + // auto last_query = na_mask.get_backward_window_end(last_key_for_block); + // NOTE(alih): NA mask formulation guarantees this condition will never + // occur! + // if (last_query < 0 || first_query >= p.num_queries) { + // return getQueryEnd(p); // nothing to compute in this column + // } + // NOTE(alih): NA mask formulation guarantees start and end are within + // range if inputs are. + // first_query = cutlass::fast_max(0, first_query); + return align_dim_down(first_query, p.query_tile_shape); + // return (first_query / kBlockSizeI) * kBlockSizeI; + } + + // Returns how many kernel blocks will write to a given block in `grad_query` + // This is usually equal to the number of key splits, but can be different + // for instance in the causal case, or varying seqlen + static CUTLASS_HOST_DEVICE int32_t + getNumParallelBlocksForQuery(Params const& p, Dim query_start) { + auto na_mask = + NAMask(p.kernel_size, p.stride, p.num_queries_post_partitioning); + + // int16_t num_key_blocks = gemm_kernel_utils::ceil_div(p.num_keys, + // kBlockSizeJ); if (p.custom_mask_type != NoCustomMask) { + // int32_t shift = p.custom_mask_type == CausalFromBottomRight + // ? p.num_keys - p.num_queries + // : 0; + // int32_t last_query_for_block = + // cutlass::fast_min(query_start + kBlockSizeI, p.num_queries) - 1; + // int32_t last_key_for_block = + // cutlass::fast_min(last_query_for_block + shift, p.num_keys - 1); + // int32_t first_key_for_block = p.window_size == 0 + // ? 0 + // : cutlass::fast_max(query_start - p.window_size + 1 + shift, 0); + + // if (p.window_size == 0) { + // num_key_blocks = last_key_for_block / kBlockSizeJ + 1; + // } else { + // num_key_blocks = (last_key_for_block / kBlockSizeJ) - + // (first_key_for_block / kBlockSizeJ) + 1; + // } + + // if (last_key_for_block < 0 || first_key_for_block >= p.num_keys) { + // num_key_blocks = 0; + // } + //} + // return cutlass::fast_min(p.num_splits_key_device(), num_key_blocks); + auto first_key_for_block = na_mask.get_window_start(query_start); + auto last_query_for_block = + fast_min( + query_start + p.query_tile_shape, p.num_queries_post_partitioning) - + 1; + auto last_key_for_block = na_mask.get_window_end_(last_query_for_block) - 1; + return (fast_min( + p.num_splits_key, + count_tiles( + first_key_for_block, last_key_for_block, p.key_tile_shape))) + .prod32(); + } + + // Returns the next block to process + static CUTLASS_HOST_DEVICE bool incrIteration( + Params const& p, + Dim query_start, + Dim key_start, + Dim& next_query, + Dim& next_key) { + auto first_query_tile = getSmallestQueryForKey(p, key_start); + auto na_mask = + NAMask(p.kernel_size, p.stride, p.num_queries_post_partitioning); + + auto last_key_for_block = + fast_min( + key_start + p.key_tile_shape, p.num_queries_post_partitioning) - + 1; + auto last_query = na_mask.get_backward_window_end(last_key_for_block); + auto last_query_align_up = align_dim_up(last_query, p.query_tile_shape); + + // next_query = query_start + kBlockSizeI; + next_query = increment_tile_return( + query_start, p.query_tile_shape, first_query_tile, last_query_align_up); + next_key = key_start; + // auto query_shift = getQueryStartShift(p); + //// Wrap around + // if (query_shift) { + // if (next_query >= p.num_queries) { + // next_query = getSmallestQueryForKey(p, key_start); + // return; + // } else if (query_start < query_shift && query_shift <= next_query) { + // // jump to next key + // } else { + // return; + // } + // } else { + // if (p.window_size > 0) { + // int32_t shift = p.custom_mask_type == CausalFromBottomRight + // ? p.num_keys - p.num_queries + // : 0; + // last key that is not masked out + // int last_key_for_block = + // cutlass::fast_min(key_start + kBlockSizeJ, p.num_keys) - 1; + // int last_query = last_key_for_block - shift + p.window_size - 1; + // if (next_query <= last_query && next_query < p.num_queries) { + // return; + // } + // NOTE(alih): NAmask determines last_query, and it won't ever be out of + // bounds as long as key_start isn't out of bounds. if next_query <= + // last_query, we don't need to check if it's out of bounds. We also don't + // need the following extra condition, because we don't want to break the + // query loop early as soon as we're done processing relevant query tiles. + // if (is_coord_less_than_or_equal_to(next_query, last_query)) { + if (is_coord_within_upper_bound(next_query, last_query) && + is_coord_within_upper_bound( + next_query, p.num_queries_post_partitioning)) { + return true; + } + // } else if (next_query < p.num_queries) { + // return; + // } + // jump to next key + //} + // Next key + // next_key = key_start + p.num_splits_key_device() * kBlockSizeJ; + increment_tile( + next_key, + p.key_tile_shape * p.num_splits_key, + /*lower_bound=*/p.split_key_device_dim() * p.key_tile_shape, + p.num_queries_post_partitioning); + next_query = getQueryStart(p, next_key); + return false; + } + + template + static CUTLASS_DEVICE void prologueQkNextIteration( + SharedStorage& shared_storage, + Params const& p, + Dim query_start, + Dim key_start, + uint8_t warp_id, + uint8_t lane_id) { + // if (query_start >= p.num_queries || key_start >= p.num_keys) { + if (!is_coord_within_upper_bound( + query_start, p.num_queries_post_partitioning) || + !is_coord_within_upper_bound( + key_start, p.num_queries_post_partitioning)) { + return; + } + auto n_rows_a = + fast_min(p.key_tile_shape, p.num_queries_post_partitioning - key_start); + auto n_rows_b = fast_min( + p.query_tile_shape, p.num_queries_post_partitioning - query_start); + // auto n_rows_a = p.num_queries_post_partitioning - key_start; + // auto n_rows_b = p.num_queries_post_partitioning - query_start; + + static constexpr bool kReloadK = + kForceReloadK || !MatmulQK::Mma::kSmemContainsEntireMat; + int thread_id = 32 * warp_id + lane_id; + typename MatmulQK::Mma::IteratorA iterator_A( + //{int32_t(p.k_strideM)}, + {p.k_strideM, n_rows_a}, + // p.key_ptr + key_start * p.k_strideM, + p.key_ptr + (key_start * p.k_strideM).sum(), + //{p.num_keys - key_start, p.head_dim}, + n_rows_a, + p.head_dim, + thread_id, + cutlass::MatrixCoord{0, 0}); + + typename MatmulQK::Mma::IteratorB iterator_B( + //{int32_t(p.q_strideM)}, + {p.q_strideM, n_rows_b}, + // p.query_ptr + query_start * p.q_strideM, + p.query_ptr + (query_start * p.q_strideM).sum(), + //{p.head_dim, p.num_queries - query_start}, + n_rows_b, + p.head_dim, + thread_id, + cutlass::MatrixCoord{0, 0}); + + MatmulQK::Mma::prologue( + shared_storage.mm_qk_k(), + shared_storage.mm_qk_q(), + iterator_A, + iterator_B, + thread_id, + p.head_dim); + } + + static CUTLASS_DEVICE void writeFragsToGmem( + SharedStorage& shared_storage, + OutputFragments& output_frags, + Params const& p, + Dim key_start, + uint8_t warp_id, + uint8_t lane_id) { + uint16_t thread_id = 32 * warp_id + lane_id; + // int32_t num_keys_in_block = + // cutlass::fast_min((int32_t)MatmulQK::Mma::Shape::kM, p.num_keys - + // key_start); + auto num_keys_in_block = + fast_min(p.key_tile_shape, p.num_queries_post_partitioning - key_start); + typename MatmulGradV::OutputTileIterator outputV_it( + // typename MatmulGradV::OutputTileIterator::Params{p.gV_strideM()}, + // p.grad_value_ptr + key_start * p.gV_strideM(), + //{num_keys_in_block, p.head_dim_value}, + typename MatmulGradV::OutputTileIterator::Params{ + p.v_strideM, num_keys_in_block}, + p.grad_value_ptr + (key_start * p.v_strideM).sum(), + num_keys_in_block, + p.head_dim_value, + thread_id); + accumulateInGmem( + shared_storage.gradV_epilogue_final(), + output_frags.gradV, + outputV_it, + true, + warp_id, + lane_id); + + typename MatmulGradK::OutputTileIterator outputK_it( + // typename MatmulGradK::OutputTileIterator::Params{p.gK_strideM()}, + // p.grad_key_ptr + key_start * p.gK_strideM(), + //{num_keys_in_block, + // false ? MatmulGradK::ThreadblockShape::kN : p.head_dim}, + typename MatmulGradK::OutputTileIterator::Params{ + p.k_strideM, num_keys_in_block}, + p.grad_key_ptr + (key_start * p.k_strideM).sum(), + num_keys_in_block, + p.head_dim, + thread_id); + accumulateInGmem( + shared_storage.gradK_epilogue_final(), + output_frags.gradK, + outputK_it, + true, + warp_id, + lane_id); + } + + template + static CUTLASS_DEVICE void accumulateInGmem( + typename MatmulT::DefaultEpilogue::SharedStorage& epilogue_smem, + typename MatmulT::Mma::FragmentC const& accum, + typename MatmulT::OutputTileIterator output_it, + bool first, + uint8_t warp_id, + uint8_t lane_id) { + using DefaultEpilogue = typename MatmulT::DefaultEpilogue; + using DefaultOutputOp = typename MatmulT::DefaultOutputOp; + using Mma = typename MatmulT::Mma; + int thread_id = 32 * warp_id + lane_id; + DISPATCH_BOOL( + first, kIsFirst, ([&]() { + static constexpr auto ScaleType = kIsFirst + ? cutlass::epilogue::thread::ScaleType::Nothing + : cutlass::epilogue::thread::ScaleType::NoBetaScaling; + using EpilogueOutputOp = + typename cutlass::epilogue::thread::LinearCombination< + typename DefaultOutputOp::ElementOutput, + DefaultOutputOp::kCount, + typename DefaultOutputOp::ElementAccumulator, + typename DefaultOutputOp::ElementCompute, + ScaleType>; + using Epilogue = + typename cutlass::epilogue::threadblock::EpiloguePipelined< + typename DefaultEpilogue::Shape, + typename Mma::Operator, + DefaultEpilogue::kPartitionsK, + typename MatmulT::OutputTileIterator, + typename DefaultEpilogue::AccumulatorFragmentIterator, + typename DefaultEpilogue::WarpTileIterator, + typename DefaultEpilogue::SharedLoadIterator, + EpilogueOutputOp, + typename DefaultEpilogue::Padding, + DefaultEpilogue::kFragmentsPerIteration, + true // IterationsUnroll + >; + EpilogueOutputOp rescale({1, 1}); + Epilogue epilogue(epilogue_smem, thread_id, warp_id, lane_id); + epilogue(rescale, output_it, accum, output_it); + })); + } + + // NOTE(alih): there's probably a bug here somewhere; + // fusing compute delta is disabled until I figure that out. + template + static CUTLASS_DEVICE void computeDelta( + Params const& p, + Dim query_start, + uint8_t warp_id, + uint8_t lane_id) { + // Each thread computes one value for Delta + // Depending on warp configuration, we might have multiple + // threads of the same warp working on the same row + using AccessType = cutlass::Array; + static_assert(kNumThreads >= kBlockSizeI, ""); + static constexpr int kNumThreadsPerLine = kNumThreads / kBlockSizeI; + int16_t thread_id = 32 * warp_id + lane_id; + + int16_t laneFirstCol = kElementsPerAccess * (lane_id % kNumThreadsPerLine); + int16_t laneRow = thread_id / kNumThreadsPerLine; + auto num_queries_in_block = fast_min( + p.query_tile_shape, p.num_queries_post_partitioning - query_start); + auto query_coord = query_start + + map_index_to_coord((int32_t)laneRow, num_queries_in_block); + bool rowPred = is_coord_within_upper_bound( + query_coord, p.num_queries_post_partitioning); + auto query_offset = (query_coord * p.o_strideM).sum(); + bool pred = rowPred; + + // on windows, previous syntax __restrict__ AccessType* + // resulted in error: "restrict" is not allowed + const AccessType* __restrict__ grad_output_ptr = + reinterpret_cast( + p.grad_output_ptr + query_offset + laneFirstCol); + const AccessType* __restrict__ output_ptr = + reinterpret_cast( + p.output_ptr + query_offset + laneFirstCol); + + static constexpr int64_t kMaxIters = + kMaxK / (kElementsPerAccess * kNumThreadsPerLine); + constexpr int kPipelineStages = 2; + accum_t delta_value = accum_t(0); + using GlobalLoad = + cutlass::arch::global_load; + AccessType frag_grad_output[kPipelineStages]; + AccessType frag_output[kPipelineStages]; + + auto loadAndIncrement = [&](int ld_pos, bool is_valid) { + frag_grad_output[ld_pos].clear(); + frag_output[ld_pos].clear(); + GlobalLoad(frag_grad_output[ld_pos], grad_output_ptr, is_valid); + GlobalLoad(frag_output[ld_pos], output_ptr, is_valid); + grad_output_ptr += kNumThreadsPerLine; + output_ptr += kNumThreadsPerLine; + }; + + CUTLASS_PRAGMA_UNROLL + for (int iter = 0; iter < kPipelineStages - 1; ++iter) { + int ld_pos = iter % kPipelineStages; + pred = pred && + (laneFirstCol + iter * kElementsPerAccess * kNumThreadsPerLine) < + p.head_dim_value; + loadAndIncrement(ld_pos, pred); + } + auto columnIteration = [&](int iter) { + // Load for next iter + int ld_pos = (iter + kPipelineStages - 1) % kPipelineStages; + pred = pred && + (laneFirstCol + + (iter + kPipelineStages - 1) * kElementsPerAccess * + kNumThreadsPerLine) < p.head_dim_value; + loadAndIncrement(ld_pos, pred); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < AccessType::kElements; ++i) { + delta_value += accum_t(frag_output[iter % kPipelineStages][i]) * + accum_t(frag_grad_output[iter % kPipelineStages][i]); + } + }; + + // If we have a small lower-bound for K, we can unroll the loop + if (kMaxK <= 256) { + CUTLASS_PRAGMA_UNROLL + for (int iter = 0; iter < kMaxIters; ++iter) { + columnIteration(iter); + } + } else { + int num_iters = + gemm_kernel_utils::ceil_div( + p.head_dim_value, kElementsPerAccess * kNumThreadsPerLine) * + (kElementsPerAccess * kNumThreadsPerLine); + for (int iter = 0; iter < num_iters; ++iter) { + columnIteration(iter); + } + } + + // Reduce between workers + static_assert( + kNumThreadsPerLine == 1 || kNumThreadsPerLine == 2 || + kNumThreadsPerLine == 4, + ""); + CUTLASS_PRAGMA_UNROLL + for (int i = 1; i < kNumThreadsPerLine; i *= 2) { + delta_value = delta_value + __shfl_xor_sync(0xffffffff, delta_value, i); + } + + // Store in gmem + if (rowPred) { + p.delta_ptr[(query_coord * p.lse_strideM).sum()] = delta_value; + } + } +}; + +} // namespace fna +} // namespace cuda +} // namespace natten diff --git a/natten/csrc/include/natten/cuda/fna/kernel_forward.h b/natten/csrc/include/natten/cuda/fna/kernel_forward.h new file mode 100644 index 00000000..02dd3be1 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna/kernel_forward.h @@ -0,0 +1,1138 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Fused Neighborhood Attention kernels are heavily based on the + * memory-efficient attention kernels from the xFormers project by Meta + * Platforms, Inc. + * + * Copyright (c) Facebook, Inc. and its affiliates + * + * BSD 3-Clause License + * + * 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 names of Facebook, Deepmind Technologies, NYU, NEC + * Laboratories America and IDIAP Research Institute 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 OWNER 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 +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include + +namespace natten { +namespace cuda { +namespace fna { + +namespace { +template +constexpr int getWarpsPerSmFw() { + return ( + Arch::kMinComputeCapability >= 80 && + !cutlass::platform::is_same::value + ? 16 + : 12); +} +static CUTLASS_DEVICE float atomicMaxFloat(float* addr, float value) { + // source: https://stackoverflow.com/a/51549250 + return (value >= 0) + ? __int_as_float(atomicMax((int*)addr, __float_as_int(value))) + : __uint_as_float(atomicMin((unsigned int*)addr, __float_as_uint(value))); +} + +} // namespace + +template < + // NA spatial dimension; can be 1 (NA1d), 2 (NA2d), or 3 (NA3d). + int NADim_, + typename CausalMask, + // The datatype of Q/K/V + typename scalar_t_, + // Architecture we are targeting (eg `cutlass::arch::Sm80`) + typename ArchTag, + // If Q/K/V are correctly aligned in memory and we can run a fast kernel + bool isAligned_, + int kQueriesPerBlock_, + int kKeysPerBlock_, + // upperbound on `max(value.shape[-1], query.shape[-1])` + int kMaxK_ = (int)cutlass::platform::numeric_limits::max()> +struct FusedNeighborhoodAttentionKernel { + static constexpr int NADim = NADim_; + static_assert(NADim >= 1 && NADim < 4, "Only 1D-3D NA are implemented."); + using Dim = typename GetDim::type; + using NAMask = NeighborhoodAttentionMask; + + static constexpr int kKeysPerBlock = kKeysPerBlock_; + static constexpr int kQueriesPerBlock = kQueriesPerBlock_; + static constexpr int kMaxK = kMaxK_; + + // static constexpr Dim QueryTileShape = MapTileSizeToNd::value; static constexpr Dim KeyTileShape = + // MapTileSizeToNd::value; + static constexpr bool kHasCausalDims = CausalMask::AnyCausalDims; + + using scalar_t = scalar_t_; + using accum_t = float; + using lse_scalar_t = float; + using output_t = scalar_t; + // Accumulator between 2 iterations + // Using `accum_t` improves perf on f16 at the cost of + // numerical errors + using output_accum_t = accum_t; + static constexpr bool kIsAligned = isAligned_; + static constexpr bool kSingleValueIteration = kMaxK <= kKeysPerBlock; + static constexpr int32_t kAlignLSE = 32; // block size of backward + static constexpr bool kIsHalf = cutlass::sizeof_bits::value == 16; + static constexpr bool kPreloadV = + ArchTag::kMinComputeCapability >= 80 && kIsHalf; + static constexpr bool kKeepOutputInRF = kSingleValueIteration; + static constexpr bool kNeedsOutputAccumulatorBuffer = !kKeepOutputInRF && + !cutlass::platform::is_same::value; + + static_assert(kQueriesPerBlock % 32 == 0, ""); + static_assert(kKeysPerBlock % 32 == 0, ""); + static constexpr int kNumWarpsPerBlock = + kQueriesPerBlock * kKeysPerBlock / (32 * 32); + static constexpr int kWarpSize = 32; + + // Launch bounds + static constexpr int kNumThreads = kWarpSize * kNumWarpsPerBlock; + static constexpr int kMinBlocksPerSm = + getWarpsPerSmFw() / kNumWarpsPerBlock; + + struct Params { + // Input tensors + scalar_t* query_ptr = nullptr; // [..., num_heads, head_dim] + scalar_t* key_ptr = nullptr; // [..., num_heads, head_dim] + scalar_t* value_ptr = nullptr; // [..., num_heads, head_dim_value] + + // Output tensors + output_t* output_ptr = + nullptr; // [num_queries_post_partitioning, num_heads, head_dim_value] + // [num_queries_post_partitioning, num_heads, head_dim_value] + output_accum_t* output_accum_ptr = nullptr; + // [num_heads, num_queries_post_partitioning] - can be null + lse_scalar_t* logsumexp_ptr = nullptr; + + // Sliding window. ignored if == 0 + Dim kernel_size; + Dim stride; + Dim dilation; + Dim cta_shape_x; + + // Scale + accum_t scale = 0.0; + + // Dimensions/strides + int32_t head_dim = 0; + int32_t head_dim_value = 0; + Dim num_queries; + Dim num_queries_post_partitioning; + + Dim query_tile_shape; + Dim key_tile_shape; + + Dim lse_strideM; + Dim q_strideM; + Dim k_strideM; + Dim v_strideM; + Dim o_strideM; + + int32_t num_heads = 0; + int32_t num_batches = 0; + + bool is_fully_block_sparse = false; + bool has_kv_padding = false; + + // Moves pointers to what we should process + // Returns "false" if there is no work to do + CUTLASS_DEVICE bool advance_to_block() { + num_queries_post_partitioning = ceil_div_dim(num_queries, dilation); + + auto batch_id = blockIdx.z; + auto bidxy = blockIdx.y; + auto dilation_size = dilation.prod32(); + auto dilation_dim_idx = (int32_t)bidxy % dilation_size; + auto dilation_idx = map_index_to_coord(dilation_dim_idx, dilation); + auto head_id = bidxy / dilation_size; + + auto q_strideH = head_dim; + auto k_strideH = head_dim; + auto v_strideH = head_dim_value; + auto o_strideH = head_dim_value; + + auto q_heads_dims = num_heads * q_strideH; + auto k_heads_dims = num_heads * k_strideH; + auto v_heads_dims = num_heads * v_strideH; + auto o_heads_dims = num_heads * o_strideH; + + auto q_stride_dilation = compute_stride(num_queries, q_heads_dims); + auto k_stride_dilation = compute_stride(num_queries, k_heads_dims); + auto v_stride_dilation = compute_stride(num_queries, v_heads_dims); + auto o_stride_dilation = compute_stride(num_queries, o_heads_dims); + + q_strideM = q_stride_dilation * dilation; + k_strideM = k_stride_dilation * dilation; + v_strideM = v_stride_dilation * dilation; + o_strideM = o_stride_dilation * dilation; + + auto q_strideB = num_queries.prod() * q_heads_dims; + auto k_strideB = num_queries.prod() * k_heads_dims; + auto v_strideB = num_queries.prod() * v_heads_dims; + auto o_strideB = num_queries.prod() * o_heads_dims; + + // This is where qkv shape is corrected: + maybe_mask_qk_tiles( + num_queries_post_partitioning, num_queries, dilation, dilation_idx); + + // NOTE: causal mask is antithetical to being fully block sparse. + // While there are cases where a causal mask has mostly dense blocks, the + // last one (at least) will always be causally masked. Fast paths for + // block sparsity in FNA are limited to masks that are guaranteed to be + // fully block sparse only. + is_fully_block_sparse = not kHasCausalDims && + fully_block_sparse( + num_queries_post_partitioning, + kernel_size, + stride, + query_tile_shape, + key_tile_shape); + + has_kv_padding = + not evenly_divides(num_queries_post_partitioning, key_tile_shape); + + cta_shape_x = + ceil_div_dim(num_queries_post_partitioning, query_tile_shape); + + const auto first_query = + map_index_to_coord((int32_t)blockIdx.x, cta_shape_x) * + query_tile_shape; + + // Advance to current batch + query_ptr += batch_id * q_strideB; + key_ptr += batch_id * k_strideB; + value_ptr += batch_id * v_strideB; + output_ptr += batch_id * o_strideB; + if (output_accum_ptr != nullptr) { + output_accum_ptr += batch_id * o_strideB; + } + + // Advance to the current batch / head / first_query + query_ptr += (first_query * q_strideM).sum() + + (dilation_idx * q_stride_dilation).sum() + head_id * q_strideH; + key_ptr += (dilation_idx * k_stride_dilation).sum() + head_id * k_strideH; + value_ptr += + (dilation_idx * v_stride_dilation).sum() + head_id * v_strideH; + output_ptr += (first_query * o_strideM).sum() + + (dilation_idx * o_stride_dilation).sum() + head_id * o_strideH; + + if (output_accum_ptr != nullptr) { + output_accum_ptr += (first_query * o_strideM).sum() + + (dilation_idx * o_stride_dilation).sum() + head_id * o_strideH; + } else { + // Accumulate directly in the destination buffer (eg for f32) + output_accum_ptr = (accum_t*)output_ptr; + } + + if (logsumexp_ptr != nullptr) { + auto lse_stride_dilation = compute_stride(num_queries, num_heads); + lse_strideM = lse_stride_dilation * dilation; + + // NOTE(alih): we don't pad for 128-bit alignment like in xFormers, + // because we have to rewrite the vector iterator in the backward + // kernel into a gather op. + // auto lse_dim = cutlass::ceil_div((int32_t)num_queries.prod32(), + // kAlignLSE) * kAlignLSE; + // lse[batch_id, query_start, head_id] + logsumexp_ptr += batch_id * num_queries.prod32() * num_heads + + ((first_query * lse_strideM).sum() + + (dilation_idx * lse_stride_dilation).sum()) + + head_id; + } + + num_batches = 0; // no longer used after + // dilation = 0; + + // Make sure the compiler knows these variables are the same on all + // the threads of the warp. + // Only worth doing if they could have been modified above. + query_ptr = gemm_kernel_utils::warp_uniform(query_ptr); + key_ptr = gemm_kernel_utils::warp_uniform(key_ptr); + value_ptr = gemm_kernel_utils::warp_uniform(value_ptr); + output_ptr = gemm_kernel_utils::warp_uniform(output_ptr); + output_accum_ptr = gemm_kernel_utils::warp_uniform(output_accum_ptr); + logsumexp_ptr = gemm_kernel_utils::warp_uniform(logsumexp_ptr); + num_heads = gemm_kernel_utils::warp_uniform(num_heads); + return true; + } + + __host__ dim3 getBlocksGrid() const { + return dim3( + ceil_div_dim(ceil_div_dim(num_queries, dilation), query_tile_shape) + .prod32(), + num_heads * dilation.prod32(), + num_batches); + } + + __host__ dim3 getThreadsGrid() const { + return dim3(kWarpSize, kNumWarpsPerBlock, 1); + } + }; + + struct MM0 { + /* + In this first matmul, we compute a block of `Q @ K.T`. + While the calculation result is still hot in registers, we update + `mi`, `m_prime`, `s_prime` in shared-memory, and then store this value + into a shared-memory ("AccumulatorSharedStorage") that is used later as + operand A for the second matmul (see MM1) + */ + using GemmType = gemm_kernel_utils::DefaultGemmType; + + using OpClass = typename GemmType::OpClass; + using DefaultConfig = + typename cutlass::gemm::device::DefaultGemmConfiguration< + OpClass, + ArchTag, + scalar_t, + scalar_t, + scalar_t, // ElementC + accum_t // ElementAccumulator + >; + static constexpr int kAlignmentA = + kIsAligned ? DefaultConfig::kAlignmentA : GemmType::kMinimumAlignment; + static constexpr int kAlignmentB = + kIsAligned ? DefaultConfig::kAlignmentB : GemmType::kMinimumAlignment; + using ThreadblockShape = cutlass::gemm:: + GemmShape; + using WarpShape = cutlass::gemm::GemmShape<32, 32, GemmType::WarpK>; + using DefaultMma = typename cutlass::gemm::threadblock::FindDefaultMma< + NADim, + scalar_t, // ElementA, + cutlass::layout::RowMajor, // LayoutA, + kAlignmentA, + scalar_t, // ElementB, + cutlass::layout::ColumnMajor, // LayoutB, + kAlignmentB, + accum_t, + cutlass::layout::RowMajor, // LayoutC, + OpClass, + ArchTag, // ArchTag + ThreadblockShape, // ThreadblockShape + WarpShape, // WarpShape + typename GemmType::InstructionShape, // InstructionShape + ArchTag::kMinComputeCapability >= 80 && kIsHalf + ? 4 + : DefaultConfig::kStages, + typename GemmType::Operator // Operator + >::DefaultMma; + using MmaCore = typename DefaultMma::MmaCore; + using IteratorA = typename DefaultMma::IteratorA; + using IteratorB = typename DefaultMma::IteratorB; + using DefaultThreadblockMma = typename DefaultMma::ThreadblockMma; + using Mma = typename cutlass::platform::conditional< + kSingleValueIteration, + typename MakeCustomMma::Mma, + DefaultThreadblockMma>::type; + using AccumLambdaIterator = typename DefaultMmaAccumLambdaIterator< + typename Mma::Operator::IteratorC, + accum_t, + kWarpSize>::Iterator; + static_assert( + MmaCore::WarpCount::kM * MmaCore::WarpCount::kN * + MmaCore::WarpCount::kK == + kNumWarpsPerBlock, + ""); + + // Epilogue to store to shared-memory in a format that we can use later for + // the second matmul + using B2bGemm = typename cutlass::gemm::threadblock::B2bGemm< + NADim, + typename Mma::Operator::IteratorC, + typename Mma::Operator, + scalar_t, + WarpShape, + ThreadblockShape>; + using AccumulatorSharedStorage = typename B2bGemm::AccumulatorSharedStorage; + }; + + struct MM1 { + /** + Second matmul: perform `attn @ V` where `attn` is the attention (not + normalized) and stored in shared memory + */ + using GemmType = gemm_kernel_utils::DefaultGemmType; + + using OpClass = typename GemmType::OpClass; + using DefaultConfig = + typename cutlass::gemm::device::DefaultGemmConfiguration< + OpClass, + ArchTag, + scalar_t, + scalar_t, + output_accum_t, // ElementC + accum_t // ElementAccumulator + >; + static constexpr int kAlignmentA = DefaultConfig::kAlignmentA; // from smem + static constexpr int kAlignmentB = + kIsAligned ? DefaultConfig::kAlignmentB : GemmType::kMinimumAlignment; + using ThreadblockShape = cutlass::gemm:: + GemmShape; + using WarpShape = cutlass::gemm::GemmShape<32, 32, GemmType::WarpK>; + using InstructionShape = typename GemmType::InstructionShape; + + using LayoutB = cutlass::layout::RowMajor; + using DefaultGemm = cutlass::gemm::kernel::DefaultGemm< + scalar_t, // ElementA, + cutlass::layout::RowMajor, // LayoutA, + kAlignmentA, + scalar_t, // ElementB, + LayoutB, // LayoutB, + kAlignmentB, + output_accum_t, + cutlass::layout::RowMajor, // LayoutC, + accum_t, + OpClass, + ArchTag, + ThreadblockShape, + WarpShape, + typename GemmType::InstructionShape, + typename DefaultConfig::EpilogueOutputOp, + void, // ThreadblockSwizzle - not used + ArchTag::kMinComputeCapability >= 80 && kIsHalf + ? 4 + : DefaultConfig::kStages, + false, // SplitKSerial + typename GemmType::Operator>; + + using WarpIteratorA = typename cutlass::gemm::threadblock:: + DefaultWarpIteratorAFromSharedMemory< + typename DefaultGemm::Mma::Policy::Operator::Shape, // WarpShape + typename DefaultGemm::Mma::Policy::Operator::InstructionShape, + typename DefaultGemm::Mma::Policy::Operator::IteratorA, + typename DefaultGemm::Mma::Policy>::WarpIterator; + using DefaultMmaFromSmem = + typename cutlass::gemm::threadblock::DefaultMmaFromSharedMemory< + NADim, + typename DefaultGemm::Mma, + MM0::AccumulatorSharedStorage::Shape::kN, // kMaxK + WarpIteratorA, + false>; // kScaleOperandA + using Mma = typename DefaultMmaFromSmem::Mma; + using IteratorB = typename Mma::IteratorB; + using WarpCount = typename Mma::WarpCount; + static_assert( + WarpCount::kM * WarpCount::kN * WarpCount::kK == kNumWarpsPerBlock, + ""); + + using DefaultEpilogue = typename DefaultGemm::Epilogue; + using OutputTileIterator = + typename cutlass::epilogue::threadblock::CustomPredicatedTileIterator< + NADim, + typename DefaultEpilogue::OutputTileIterator::ThreadMap, + output_t>; + using OutputTileIteratorAccum = + typename cutlass::epilogue::threadblock::CustomPredicatedTileIterator< + NADim, + typename DefaultEpilogue::OutputTileIterator::ThreadMap, + output_accum_t>; + }; + + static constexpr int64_t kAlignmentQ = MM0::kAlignmentA; + static constexpr int64_t kAlignmentK = MM0::kAlignmentB; + static constexpr int64_t kAlignmentV = 1; + + // Shared storage - depends on kernel params + struct ScalingCoefs { + cutlass::Array m_prime; + cutlass::Array s_prime; + cutlass::Array mi; + cutlass::Array out_rescale; + cutlass::Array + addition_storage; + }; + + struct SharedStorageEpilogueAtEnd : ScalingCoefs { + struct SharedStorageAfterMM0 { + // Everything here might be overwritten during MM0 + union { + typename MM0::AccumulatorSharedStorage si; + }; + typename MM1::Mma::SharedStorage mm1; + }; + + union { + typename MM0::Mma::SharedStorage mm0; + SharedStorageAfterMM0 after_mm0; + typename MM1::DefaultEpilogue::SharedStorage epilogue; + }; + + CUTLASS_DEVICE typename MM1::DefaultEpilogue::SharedStorage& + epilogue_shared_storage() { + return epilogue; + } + }; + + struct SharedStorageEpilogueInLoop : ScalingCoefs { + struct SharedStorageAfterMM0 { + // Everything here might be overwritten during MM0 + union { + typename MM0::AccumulatorSharedStorage si; + }; + typename MM1::Mma::SharedStorage mm1; + typename MM1::DefaultEpilogue::SharedStorage epilogue; + }; + + union { + typename MM0::Mma::SharedStorage mm0; + SharedStorageAfterMM0 after_mm0; + }; + + CUTLASS_DEVICE typename MM1::DefaultEpilogue::SharedStorage& + epilogue_shared_storage() { + return after_mm0.epilogue; + } + }; + + using SharedStorage = typename cutlass::platform::conditional< + kSingleValueIteration || kKeepOutputInRF, + SharedStorageEpilogueAtEnd, + SharedStorageEpilogueInLoop>::type; + + static bool __host__ check_supported(Params const& p) { + CHECK_ALIGNED_PTR(p.query_ptr, kAlignmentQ); + CHECK_ALIGNED_PTR(p.key_ptr, kAlignmentK); + CHECK_ALIGNED_PTR(p.value_ptr, kAlignmentV); + NATTEN_CHECK( + p.head_dim % kAlignmentQ == 0, + "query is not correctly aligned (strideM)"); + NATTEN_CHECK( + p.head_dim % kAlignmentK == 0, + "key is not correctly aligned (strideM)"); + NATTEN_CHECK( + p.head_dim_value % kAlignmentV == 0, + "value is not correctly aligned (strideM)"); + return true; + } + + static void CUTLASS_DEVICE attention_kernel(Params& p) { + // In this block, we will only ever: + // - read query[first_query:last_query, :] + // - write to output[first_query:last_query, :] + + extern __shared__ char smem_buffer[]; + SharedStorage& shared_storage = *((SharedStorage*)smem_buffer); + auto& m_prime = shared_storage.m_prime; + auto& s_prime = shared_storage.s_prime; + auto& mi = shared_storage.mi; + auto& out_rescale = shared_storage.out_rescale; + auto first_query = map_index_to_coord((int32_t)blockIdx.x, p.cta_shape_x) * + p.query_tile_shape; + auto problem_size_0_m = fast_min( + p.query_tile_shape, p.num_queries_post_partitioning - first_query); + auto last_query = first_query + problem_size_0_m - 1; + auto problem_size_0_m_int = problem_size_0_m.prod32(); + + static_assert(kQueriesPerBlock < kNumWarpsPerBlock * kWarpSize, ""); + if (thread_id() < kQueriesPerBlock) { + s_prime[thread_id()] = accum_t(0); + out_rescale[thread_id()] = accum_t(1.0); + m_prime[thread_id()] = + -cutlass::platform::numeric_limits::infinity(); + mi[thread_id()] = -cutlass::platform::numeric_limits::infinity(); + } + typename MM1::Mma::FragmentC accum_o; + accum_o.clear(); + + auto createOutputIter = [&](int col) -> typename MM1::OutputTileIterator { + using OutputTileIterator = typename MM1::OutputTileIterator; + return OutputTileIterator( + typename OutputTileIterator::Params{p.o_strideM, problem_size_0_m}, + p.output_ptr, + problem_size_0_m, + p.head_dim_value, + thread_id(), + {0, col}); + }; + + auto createOutputAccumIter = [&](int col) -> + typename MM1::OutputTileIteratorAccum { + using OutputTileIteratorAccum = typename MM1::OutputTileIteratorAccum; + return OutputTileIteratorAccum( + typename OutputTileIteratorAccum::Params{ + p.o_strideM, problem_size_0_m}, + p.output_accum_ptr, + problem_size_0_m, + p.head_dim_value, + thread_id(), + {0, col}); + }; + + int32_t const& problem_size_0_k_int = p.head_dim; + int32_t const& problem_size_1_n_int = p.head_dim_value; + auto gemm_k_iterations = + (problem_size_0_k_int + MM0::Mma::Shape::kK - 1) / MM0::Mma::Shape::kK; + const int64_t nBlockN = kSingleValueIteration + ? 1 + : gemm_kernel_utils::ceil_div( + (int64_t)problem_size_1_n_int, + int64_t(MM1::ThreadblockShape::kN)); + + auto na_mask = + NAMask(p.kernel_size, p.stride, p.num_queries_post_partitioning); + + auto first_key = na_mask.get_window_start(first_query); + auto last_key = na_mask.get_window_end_(last_query); + + auto my_warp_id = gemm_kernel_utils::warp_uniform(warp_id()); + auto my_lane_id = lane_id(); + + typename MM0::Mma::Operator::IteratorC::TensorCoord iteratorC_tile_offset = + {/*(tb_tile_offset.m() * MM0::Mma::WarpCount::kM) +*/ + (my_warp_id % MM0::Mma::WarpCount::kM), + /*(tb_tile_offset.n() * MM0::Mma::WarpCount::kN) +*/ + (my_warp_id / MM0::Mma::WarpCount::kM)}; + + auto lane_offset = MM0::AccumLambdaIterator::get_lane_offset( + my_lane_id, my_warp_id, iteratorC_tile_offset); + + int warp_idx_mn_0 = my_warp_id % + (MM0::Mma::Base::WarpCount::kM * MM0::Mma::Base::WarpCount::kN); + auto output_tile_coords = cutlass::MatrixCoord{ + warp_idx_mn_0 % MM0::Mma::Base::WarpCount::kM, + warp_idx_mn_0 / MM0::Mma::Base::WarpCount::kM}; + + // Iterate through keys + for (auto iter_key_start = first_key; + is_coord_within_upper_bound(iter_key_start, last_key); + increment_tile( + iter_key_start, p.key_tile_shape, first_key, last_key)) { + const auto problem_size_0_n = + fast_min(p.key_tile_shape, last_key - iter_key_start); + auto const& problem_size_1_k = problem_size_0_n; + auto problem_size_0_n_int = problem_size_0_n.prod32(); + auto& gemm_1_num_iterations = problem_size_0_n_int; + + auto& last_kv_col = gemm_1_num_iterations; + bool is_first_kv_iter = iter_key_start == first_key; + bool is_last_kv_iter = + is_tile_last(iter_key_start, p.key_tile_shape, last_key); + + auto prologueV = [&](int blockN) { + typename MM1::Mma::IteratorB iterator_V( + typename MM1::IteratorB::Params{p.v_strideM, problem_size_1_k}, + p.value_ptr + (iter_key_start * p.v_strideM).sum(), + problem_size_1_k, + problem_size_1_n_int, + thread_id(), + cutlass::MatrixCoord{0, blockN * MM1::Mma::Shape::kN}); + MM1::Mma::prologue( + shared_storage.after_mm0.mm1, + iterator_V, + thread_id(), + gemm_1_num_iterations); + }; + + __syncthreads(); // Need to have shared memory initialized, and `m_prime` + // updated from end of prev iter + // + // MATMUL: Q.K_t + // + // Computes the block-matrix product of: + // (a) query[first_query:last_query, :] + // with + // (b) key[iter_key_start:iter_key_start + kKeysPerBlock] + // and stores that into `shared_storage.si` + // + + // Construct iterators to A and B operands + typename MM0::IteratorA iterator_A( + typename MM0::IteratorA::Params(p.q_strideM, problem_size_0_m), + p.query_ptr, + problem_size_0_m, + problem_size_0_k_int, + thread_id()); + + typename MM0::IteratorB iterator_B( + typename MM0::IteratorB::Params(p.k_strideM, problem_size_0_n), + p.key_ptr + (iter_key_start * p.k_strideM).sum(), + problem_size_0_n, + problem_size_0_k_int, + thread_id()); + + // Construct thread-scoped matrix multiply + typename MM0::Mma mma( + shared_storage.mm0, thread_id(), my_warp_id, my_lane_id); + + typename MM0::Mma::FragmentC accum; + + accum.clear(); + + // Compute threadblock-scoped matrix multiply-add + mma(gemm_k_iterations, accum, iterator_A, iterator_B, accum); + __syncthreads(); + + if (kPreloadV) { + prologueV(0); + } else { + MM1::Mma::drain_cp_asyncs(); + } + + if (not p.is_fully_block_sparse) { + // Neighborhood Attention masking + Dim first_col, key_bound, row_idx; + MM0::AccumLambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) { + row_idx = map_index_to_coord((int32_t)accum_m, problem_size_0_m) + + first_query; + first_col = na_mask.get_window_start(row_idx); + key_bound = + na_mask.get_window_end(row_idx, first_col) - first_col; + }, + [&](int accum_m, int accum_n, int idx) { + auto col = + map_index_to_coord((int32_t)accum_n, problem_size_0_n) + + iter_key_start; + if (!is_coord_within_bounds_nn(col - first_col, key_bound)) { + accum[idx] = + -cutlass::platform::numeric_limits::infinity(); + } + }, + [&](int accum_m) {}); + } else if (p.has_kv_padding) { + // Residual masking + Dim first_col, key_bound, row_idx; + MM0::AccumLambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) {}, + [&](int accum_m, int accum_n, int idx) { + auto col = + map_index_to_coord((int32_t)accum_n, problem_size_0_n) + + iter_key_start; + if (!is_coord_less_than_or_equal_to( + col, p.num_queries_post_partitioning)) { + accum[idx] = + -cutlass::platform::numeric_limits::infinity(); + } + }, + [&](int accum_m) {}); + } + + // Update `mi` from accum stored in registers + // Also does accum[i] <- exp(accum[i] - mi) + iterative_softmax( + accum_o, + accum, + mi, + m_prime, + s_prime, + out_rescale, + shared_storage.addition_storage, + my_lane_id, + thread_id(), + my_warp_id, + last_kv_col, + is_first_kv_iter, + iteratorC_tile_offset, + p.scale); + + // Output results to shared-memory + + MM0::B2bGemm::accumToSmem( + shared_storage.after_mm0.si, accum, my_lane_id, output_tile_coords); + + __syncthreads(); + + // + // MATMUL: Attn . V + // Run the matmul `attn @ V` for a block of attn and V. + // `attn` is read from shared memory (in `shared_storage_si`) + // `V` is read from global memory (with iterator_B) + // + + for (int blockN = 0; blockN < nBlockN; ++blockN) { + int gemm_k_iterations = + (gemm_1_num_iterations + MM1::Mma::Shape::kK - 1) / + MM1::Mma::Shape::kK; + + // Compute threadblock-scoped matrix multiply-add and store it in accum + // (in registers) + if (!kPreloadV) { + __syncthreads(); // we share shmem between mma and epilogue + } + + typename MM1::Mma::IteratorB iterator_V( + typename MM1::IteratorB::Params{p.v_strideM, problem_size_1_k}, + p.value_ptr + (iter_key_start * p.v_strideM).sum(), + problem_size_1_k, + problem_size_1_n_int, + thread_id(), + cutlass::MatrixCoord{0, blockN * MM1::Mma::Shape::kN}); + typename MM1::Mma mma_pv( + // operand A: Pij_dropped in shared memory + shared_storage.after_mm0.si.accum_ref(), + // operand B: shared memory staging area for Vj, which is loaded + // from global memory + shared_storage.after_mm0.mm1.operand_B_ref(), + (int)thread_id(), + (int)my_warp_id, + (int)my_lane_id); + mma_pv.set_prologue_done(kPreloadV); + if (!kKeepOutputInRF) { + accum_o.clear(); + } + mma_pv(gemm_k_iterations, accum_o, iterator_V, accum_o); + __syncthreads(); + + if (kPreloadV && !kSingleValueIteration && blockN + 1 < nBlockN) { + prologueV(blockN + 1); + } + + if (!kKeepOutputInRF) { + MM1::Mma::drain_cp_asyncs(); + DISPATCH_BOOL( + is_first_kv_iter, kIsFirst, ([&] { + DISPATCH_BOOL( + is_last_kv_iter, kIsLast, ([&] { + using DefaultEpilogue = typename MM1::DefaultEpilogue; + using DefaultOp = + typename MM1::DefaultConfig::EpilogueOutputOp; + using ElementCompute = typename DefaultOp::ElementCompute; + using EpilogueOutputOp = typename cutlass::epilogue:: + thread::MemoryEfficientAttentionNormalize< + typename cutlass::platform::conditional< + kIsLast, + output_t, + output_accum_t>::type, + output_accum_t, + DefaultOp::kCount, + typename DefaultOp::ElementAccumulator, + ElementCompute, + kIsFirst, + kIsLast, + cutlass::Array>; + using Epilogue = typename cutlass::epilogue::threadblock:: + EpiloguePipelined< + typename DefaultEpilogue::Shape, + typename MM1::Mma::Operator, + DefaultEpilogue::kPartitionsK, + typename cutlass::platform::conditional< + kIsLast, + typename MM1::OutputTileIterator, + typename MM1::OutputTileIteratorAccum>::type, + typename DefaultEpilogue:: + AccumulatorFragmentIterator, + typename DefaultEpilogue::WarpTileIterator, + typename DefaultEpilogue::SharedLoadIterator, + EpilogueOutputOp, + typename DefaultEpilogue::Padding, + DefaultEpilogue::kFragmentsPerIteration, + true, // IterationsUnroll + typename MM1::OutputTileIteratorAccum // Read + // iterator + >; + + int col = blockN * MM1::Mma::Shape::kN; + auto source_iter = createOutputAccumIter(col); + auto dest_iter = gemm_kernel_utils::call_conditional< + kIsLast, + decltype(createOutputIter), + decltype(createOutputAccumIter)>:: + apply(createOutputIter, createOutputAccumIter, col); + EpilogueOutputOp rescale(s_prime, out_rescale); + Epilogue epilogue( + shared_storage.epilogue_shared_storage(), + thread_id(), + my_warp_id, + my_lane_id); + epilogue(rescale, dest_iter, accum_o, source_iter); + })); + })); + if (!kSingleValueIteration) { + __syncthreads(); + } + } + } + __syncthreads(); // we modify `m_prime` after + } + + if (kKeepOutputInRF) { + constexpr bool kIsFirst = true; + constexpr bool kIsLast = true; + using DefaultEpilogue = typename MM1::DefaultEpilogue; + using DefaultOp = typename MM1::DefaultConfig::EpilogueOutputOp; + using ElementCompute = typename DefaultOp::ElementCompute; + using EpilogueOutputOp = + typename cutlass::epilogue::thread::MemoryEfficientAttentionNormalize< + output_t, // output + output_accum_t, // source + DefaultOp::kCount, + typename DefaultOp::ElementAccumulator, // accum + output_accum_t, // compute + kIsFirst, + kIsLast, + cutlass::Array>; + using Epilogue = + typename cutlass::epilogue::threadblock::EpiloguePipelined< + typename DefaultEpilogue::Shape, + typename MM1::Mma::Operator, + DefaultEpilogue::kPartitionsK, + typename MM1::OutputTileIterator, // destination + typename DefaultEpilogue::AccumulatorFragmentIterator, + typename DefaultEpilogue::WarpTileIterator, + typename DefaultEpilogue::SharedLoadIterator, + EpilogueOutputOp, + typename DefaultEpilogue::Padding, + DefaultEpilogue::kFragmentsPerIteration, + true, // IterationsUnroll + typename MM1::OutputTileIteratorAccum // source tile + >; + auto dest_iter = createOutputIter(0); + EpilogueOutputOp rescale(s_prime, out_rescale); + Epilogue epilogue( + shared_storage.epilogue_shared_storage(), + thread_id(), + warp_id(), + lane_id()); + MM1::Mma::drain_cp_asyncs(); + epilogue(rescale, dest_iter, accum_o); + } + + // 7. Calculate logsumexp + // To make the backward easier, we pad logsumexp with `inf` + // this avoids a few bound checks, and is not more expensive during fwd + static_assert(kQueriesPerBlock < kNumWarpsPerBlock * kWarpSize, ""); + if (p.logsumexp_ptr && thread_id() < kQueriesPerBlock) { + // auto lse_dim = cutlass::ceil_div((int32_t)p.num_queries.prod32(), + // kAlignLSE) * kAlignLSE; + constexpr float kLog2e = 1.4426950408889634074; // log_2(e) = M_LOG2E + auto query_idx = + map_index_to_coord((int32_t)thread_id(), problem_size_0_m); + auto query_offset = (query_idx * p.lse_strideM).sum(); + if (is_coord_within_upper_bound(query_idx, problem_size_0_m)) { + p.logsumexp_ptr[query_offset] = accum_t(mi[thread_id()] / kLog2e) + + cutlass::fast_log(accum_t(s_prime[thread_id()])); + //} else if (query_offset < lse_dim) { + // p.logsumexp_ptr[query_offset] = + // cutlass::platform::numeric_limits::infinity(); + } + } + } + + template + CUTLASS_DEVICE static void iterative_softmax( + typename WarpIteratorC::Fragment& frag_o, // output so far + typename WarpIteratorC::Fragment& frag, + cutlass::Array& mi, + cutlass::Array& m_prime, + cutlass::Array& s_prime, + cutlass::Array& out_rescale, + cutlass::Array& + addition_storage, + int8_t lane_id, + int8_t thread_id, + int8_t warp_id, + int max_col, + bool is_first, + typename WarpIteratorC::TensorCoord const& tile_offset, + float scaling) { + /* Iterates on the accumulator and corresponding position on result matrix + + (1) Update `mi[r]` to the max value of the row `r` + (2) In a second iteration do the following: + (a) accum <- exp(accum - mi) + (b) m_prime <- exp(m_prime - mi) + (c) s_prime <- s_prime * m_prime + sum(accum) + + All of this is done on registers, before we store all of this + on shared memory for the next matmul with Value. + */ + using Fragment = typename WarpIteratorC::Fragment; + using LambdaIterator = typename DefaultMmaAccumLambdaIterator< + WarpIteratorC, + accum_t, + kWarpSize>::Iterator; + // Convert to `accum_t` (rather than double) + constexpr float kLog2e = 1.4426950408889634074; // log_2(e) = M_LOG2E + + static_assert(kQueriesPerBlock % kNumWarpsPerBlock == 0, ""); + static constexpr int kLinesPerWarp = kQueriesPerBlock / kNumWarpsPerBlock; + + frag = cutlass::multiplies()(scaling * kLog2e, frag); + + auto lane_offset = + LambdaIterator::get_lane_offset(lane_id, warp_id, tile_offset); + + // First update `mi` to the max per-row + { + accum_t max; + LambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) { + max = -cutlass::platform::numeric_limits::infinity(); + }, + [&](int accum_m, int accum_n, int idx) { + if (accum_n < max_col) { + max = cutlass::fast_max(max, frag[idx]); + } + }, + [&](int accum_m) { + // Having 4x atomicMax seems faster than reduce within warp + // first... + atomicMaxFloat(&mi[accum_m], max); + }); + } + + // Make sure we all share the update values for `mi` + __syncthreads(); + + // Doing this `exp` is quite expensive. Let's + // split it across the warps + bool restore_mi_to_minus_inf = false; + if (lane_id < kLinesPerWarp) { + int id = warp_id * kLinesPerWarp + lane_id; + auto m_prime_id = m_prime[id]; + auto mi_id = mi[id]; + bool changed = m_prime_id < mi_id; // `false` if both are -inf + if (changed) { + auto m_prime_exp = exp2f(m_prime_id - mi_id); + out_rescale[id] = m_prime_exp; + s_prime[id] *= m_prime_exp; + } else { + out_rescale[id] = 1.0f; + } + } + __syncthreads(); // Update output fragments + if (kKeepOutputInRF && !is_first) { + accum_t line_rescale; + LambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) { line_rescale = out_rescale[accum_m]; }, + [&](int accum_m, int accum_n, int idx) { + frag_o[idx] = frag_o[idx] * line_rescale; + }, + [&](int accum_m) {}); + } + // Update accum_m, accum_n, ... + { + accum_t mi_row, total_row; + bool mask; + LambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) { + mi_row = mi[accum_m]; + mask = mi_row == + -cutlass::platform::numeric_limits::infinity(); + }, + [&](int accum_m, int accum_n, int idx) { + frag[idx] = ((!mask) && (accum_n < max_col)) + ? exp2f(frag[idx] - mi_row) + : accum_t(0.0); + }, + [&](int accum_m) {}); + LambdaIterator::iterateRows( + lane_offset, + [&](int accum_m) { total_row = 0.0; }, + [&](int accum_m, int accum_n, int idx) { total_row += frag[idx]; }, + [&](int accum_m) { + if (LambdaIterator::reduceSameRow( + lane_id, total_row, [](accum_t a, accum_t b) { + return a + b; + })) { + // NOTE: we could atomically add `total_row` to `s_prime`, but + // it's faster (and deterministic) to avoid atomics here + addition_storage + [accum_m + kQueriesPerBlock * tile_offset.column()] = + total_row; + } + }); + } + __syncthreads(); + if (lane_id < kLinesPerWarp) { + int id = warp_id * kLinesPerWarp + lane_id; + accum_t total_row = s_prime[id]; + if (restore_mi_to_minus_inf) { + // Restore `mi`, see above when we set `restore_mi_to_minus_inf=true` + mi[id] = -cutlass::platform::numeric_limits::infinity(); + } else { + m_prime[id] = mi[id]; + } + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < MM0::MmaCore::WarpCount::kN; ++i) { + total_row += addition_storage[id + kQueriesPerBlock * i]; + } + s_prime[id] = total_row; + } + } + + static CUTLASS_DEVICE int8_t lane_id() { + return threadIdx.x; + } + static CUTLASS_DEVICE int8_t warp_id() { + return threadIdx.y; + } + static CUTLASS_DEVICE int16_t thread_id() { + return threadIdx.x + threadIdx.y * blockDim.x; + } +}; + +} // namespace fna +} // namespace cuda +} // namespace natten diff --git a/natten/csrc/include/natten/cuda/fna/na_utils.cuh b/natten/csrc/include/natten/cuda/fna/na_utils.cuh new file mode 100644 index 00000000..95d38fe3 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna/na_utils.cuh @@ -0,0 +1,1725 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ +/*! \file + \brief Fused NA (FNA) utils. + Holds data structures and helper functions that are templated by + spatial rank (1d, 2d, 3d), and provide (almost) generic interfaces + for the kernel. +*/ +#pragma once + +#include +#include + +#include + +namespace natten { +namespace cuda { +namespace fna { + +struct NA1dDim { + int32_t x; + + // Default ctor + CUTLASS_HOST_DEVICE constexpr NA1dDim(int32_t x_) : x(x_) {} + + CUTLASS_HOST_DEVICE NA1dDim() : x(0) {} + + CUTLASS_HOST_DEVICE void operator=(const NA1dDim& other) { + x = other.x; + } + + // Operators + CUTLASS_HOST_DEVICE NA1dDim operator+(NA1dDim other) const { + return NA1dDim(x + other.x); + } + + CUTLASS_HOST_DEVICE NA1dDim operator*(NA1dDim other) const { + return NA1dDim(x * other.x); + } + + CUTLASS_HOST_DEVICE NA1dDim operator-(NA1dDim other) const { + return NA1dDim(x - other.x); + } + + CUTLASS_HOST_DEVICE NA1dDim operator/(NA1dDim other) const { + return NA1dDim(x / other.x); + } + + CUTLASS_HOST_DEVICE NA1dDim operator+(int32_t other) const { + return NA1dDim(x + other); + } + + CUTLASS_HOST_DEVICE NA1dDim operator*(int32_t other) const { + return NA1dDim(x * other); + } + + CUTLASS_HOST_DEVICE NA1dDim operator-(int32_t other) const { + return NA1dDim(x - other); + } + + CUTLASS_HOST_DEVICE NA1dDim operator/(int32_t other) const { + return NA1dDim(x / other); + } + + // Cmp + // CUTLASS_HOST_DEVICE bool operator<=(NA1dDim other) const { + // return x <= other.x; + // } + + // CUTLASS_HOST_DEVICE bool operator<(NA1dDim other) const { + // return x < other.x; + // } + + // CUTLASS_HOST_DEVICE bool operator>(NA1dDim other) const { + // return x > other.x; + // } + + // CUTLASS_HOST_DEVICE bool operator>=(NA1dDim other) const { + // return x >= other.x; + // } + + CUTLASS_HOST_DEVICE bool operator==(NA1dDim other) const { + return x == other.x; + } + + // This method is used to determine whether or not the current + // coordinate exceeds the spatial extent. This is used in + // the NA mask (`get_window_start`). + // It should NOT overload the >= operator. + CUTLASS_HOST_DEVICE NA1dDim is_overflowing(NA1dDim extent) const { + return NA1dDim(x >= extent.x); + } + + // CUTLASS_HOST_DEVICE bool operator<=(int32_t other) const { + // return x <= other; + //} + + // CUTLASS_HOST_DEVICE bool operator<(int32_t other) const { + // return x < other; + //} + + // CUTLASS_HOST_DEVICE bool operator>(int32_t other) const { + // return x > other; + //} + + // CUTLASS_HOST_DEVICE bool operator>=(int32_t other) const { + // return x >= other; + //} + + // CUTLASS_HOST_DEVICE bool operator==(int32_t other) const { + // return x == other; + //} + + // Reduce + CUTLASS_HOST_DEVICE int64_t sum() const { + return (int64_t)x; + } + + CUTLASS_HOST_DEVICE int64_t prod() const { + return (int64_t)x; + } + + CUTLASS_HOST_DEVICE int32_t prod32() const { + return x; + } +}; + +struct NA2dDim { + int32_t x; + int32_t y; + + // Default ctor + CUTLASS_HOST_DEVICE constexpr NA2dDim(int32_t x_, int32_t y_) + : x(x_), y(y_) {} + + CUTLASS_HOST_DEVICE NA2dDim() : x(0), y(0) {} + + CUTLASS_HOST_DEVICE void operator=(const NA2dDim& other) { + x = other.x; + y = other.y; + } + + // CUTLASS_HOST_DEVICE NA2dDim(int32_t n): x(n), y(n) {} + + // Operators + CUTLASS_HOST_DEVICE NA2dDim operator+(NA2dDim other) const { + return NA2dDim(x + other.x, y + other.y); + } + + CUTLASS_HOST_DEVICE NA2dDim operator*(NA2dDim other) const { + return NA2dDim(x * other.x, y * other.y); + } + + CUTLASS_HOST_DEVICE NA2dDim operator-(NA2dDim other) const { + return NA2dDim(x - other.x, y - other.y); + } + + CUTLASS_HOST_DEVICE NA2dDim operator/(NA2dDim other) const { + return NA2dDim(x / other.x, y / other.y); + } + + CUTLASS_HOST_DEVICE NA2dDim operator+(int32_t other) const { + return NA2dDim(x + other, y + other); + } + + CUTLASS_HOST_DEVICE NA2dDim operator*(int32_t other) const { + return NA2dDim(x * other, y * other); + } + + CUTLASS_HOST_DEVICE NA2dDim operator-(int32_t other) const { + return NA2dDim(x - other, y - other); + } + + CUTLASS_HOST_DEVICE NA2dDim operator/(int32_t other) const { + return NA2dDim(x / other, y / other); + } + + // Cmp + // CUTLASS_HOST_DEVICE bool operator<=(NA2dDim other) const { + // return x <= other.x && y <= other.y; + //} + + // CUTLASS_HOST_DEVICE bool operator<(NA2dDim other) const { + // return x < other.x && y < other.y; + //} + + // CUTLASS_HOST_DEVICE bool operator>(NA2dDim other) const { + // return x > other.x && y > other.y; + //} + + // CUTLASS_HOST_DEVICE bool operator>=(NA2dDim other) const { + // return x >= other.x && y >= other.y; + //} + + CUTLASS_HOST_DEVICE bool operator==(NA2dDim other) const { + return x == other.x && y == other.y; + } + + // This method is used to determine whether or not the current + // coordinate exceeds the spatial extent. This is used in + // the NA mask (`get_window_start`). + // It should NOT overload the >= operator. + CUTLASS_HOST_DEVICE NA2dDim is_overflowing(NA2dDim extent) const { + return NA2dDim(x >= extent.x, y >= extent.y); + } + + // CUTLASS_HOST_DEVICE bool operator<=(int32_t other) const { + // return x <= other && y <= other; + //} + + // CUTLASS_HOST_DEVICE bool operator<(int32_t other) const { + // return x < other && y < other; + //} + + // CUTLASS_HOST_DEVICE bool operator>(int32_t other) const { + // return x > other && y > other; + //} + + // CUTLASS_HOST_DEVICE bool operator>=(int32_t other) const { + // return x >= other && y >= other; + //} + + // CUTLASS_HOST_DEVICE bool operator==(int32_t other) const { + // return x == other && y == other; + //} + + // Reduce + CUTLASS_HOST_DEVICE int64_t sum() const { + return (int64_t)(x + y); + } + + CUTLASS_HOST_DEVICE int64_t prod() const { + return (int64_t)(x * y); + } + + CUTLASS_HOST_DEVICE int32_t prod32() const { + return x * y; + } +}; + +struct NA3dDim { + int32_t x; + int32_t y; + int32_t z; + + // Default ctor + CUTLASS_HOST_DEVICE constexpr NA3dDim(int32_t x_, int32_t y_, int32_t z_) + : x(x_), y(y_), z(z_) {} + + CUTLASS_HOST_DEVICE NA3dDim() : x(0), y(0), z(0) {} + + CUTLASS_HOST_DEVICE void operator=(const NA3dDim& other) { + x = other.x; + y = other.y; + z = other.z; + } + + // CUTLASS_HOST_DEVICE NA3dDim(int32_t n): x(n), y(n), z(n) {} + + // Operators + CUTLASS_HOST_DEVICE NA3dDim operator+(NA3dDim other) const { + return NA3dDim(x + other.x, y + other.y, z + other.z); + } + + CUTLASS_HOST_DEVICE NA3dDim operator*(NA3dDim other) const { + return NA3dDim(x * other.x, y * other.y, z * other.z); + } + + CUTLASS_HOST_DEVICE NA3dDim operator-(NA3dDim other) const { + return NA3dDim(x - other.x, y - other.y, z - other.z); + } + + CUTLASS_HOST_DEVICE NA3dDim operator/(NA3dDim other) const { + return NA3dDim(x / other.x, y / other.y, z / other.z); + } + + CUTLASS_HOST_DEVICE NA3dDim operator+(int32_t other) const { + return NA3dDim(x + other, y + other, z + other); + } + + CUTLASS_HOST_DEVICE NA3dDim operator*(int32_t other) const { + return NA3dDim(x * other, y * other, z * other); + } + + CUTLASS_HOST_DEVICE NA3dDim operator-(int32_t other) const { + return NA3dDim(x - other, y - other, z - other); + } + + CUTLASS_HOST_DEVICE NA3dDim operator/(int32_t other) const { + return NA3dDim(x / other, y / other, z / other); + } + + // Cmp + // CUTLASS_HOST_DEVICE bool operator<=(NA3dDim other) const { + // return x <= other.x && y <= other.y && z <= other.z; + //} + + // CUTLASS_HOST_DEVICE bool operator<(NA3dDim other) const { + // return x < other.x && y < other.y && z < other.z; + //} + + // CUTLASS_HOST_DEVICE bool operator>(NA3dDim other) const { + // return x > other.x && y > other.y && z > other.z; + //} + + // CUTLASS_HOST_DEVICE bool operator>=(NA3dDim other) const { + // return x >= other.x && y >= other.y && z >= other.z; + //} + + CUTLASS_HOST_DEVICE bool operator==(NA3dDim other) const { + return x == other.x && y == other.y && z == other.z; + } + + // This method is used to determine whether or not the current + // coordinate exceeds the spatial extent. This is used in + // the NA mask (`get_window_start`). + // It should NOT overload the >= operator. + CUTLASS_HOST_DEVICE NA3dDim is_overflowing(NA3dDim extent) const { + return NA3dDim(x >= extent.x, y >= extent.y, z >= extent.z); + } + + // CUTLASS_HOST_DEVICE bool operator<=(int32_t other) const { + // return x <= other && y <= other && z <= other; + //} + + // CUTLASS_HOST_DEVICE bool operator<(int32_t other) const { + // return x < other && y < other && z < other; + //} + + // CUTLASS_HOST_DEVICE bool operator>(int32_t other) const { + // return x > other && y > other && z > other; + //} + + // CUTLASS_HOST_DEVICE bool operator>=(int32_t other) const { + // return x >= other && y >= other && z >= other; + //} + + // CUTLASS_HOST_DEVICE bool operator==(int32_t other) const { + // return x == other && y == other && z == other; + //} + + // Reduce + CUTLASS_HOST_DEVICE int64_t sum() const { + return (int64_t)(x + y + z); + } + + CUTLASS_HOST_DEVICE int64_t prod() const { + return (int64_t)(x * y * z); + } + + CUTLASS_HOST_DEVICE int32_t prod32() const { + return x * y * z; + } +}; + +template +struct GetDim; + +template <> +struct GetDim<1> { + using type = NA1dDim; +}; + +template <> +struct GetDim<2> { + using type = NA2dDim; +}; + +template <> +struct GetDim<3> { + using type = NA3dDim; +}; + +// data structure conversions +template +Dim tuple_to_na_dim(TupleType v); + +template <> +inline NA1dDim tuple_to_na_dim(std::tuple v) { + return NA1dDim(std::get<0>(v)); +} + +template <> +inline NA2dDim tuple_to_na_dim(std::tuple v) { + return NA2dDim(std::get<0>(v), std::get<1>(v)); +} + +template <> +inline NA3dDim tuple_to_na_dim(std::tuple v) { + return NA3dDim(std::get<0>(v), std::get<1>(v), std::get<2>(v)); +} + +//// Factories +// template +// struct MakeDim; +// +// template +// struct MakeDim { +// static constexpr NA1dDim value = {Value}; +//}; +// +// template +// struct MakeDim { +// static constexpr NA2dDim value = {Value, Value}; +//}; +// +// template +// struct MakeDim { +// static constexpr NA3dDim value = {Value, Value, Value}; +//}; + +//// Tile mappings +// template +// struct MapTileSizeToNd; +// +//// Identity mapping for all 1D problems. +// template +// struct MapTileSizeToNd { +// static constexpr NA1dDim value = TileSize; +//}; +// +//// 2D problems +// template <> +// struct MapTileSizeToNd<32, 2> { +// static constexpr NA2dDim value = {2, 16}; +//}; +// +// template <> +// struct MapTileSizeToNd<64, 2> { +// static constexpr NA2dDim value = {4, 16}; +//}; +// +// template <> +// struct MapTileSizeToNd<128, 2> { +// static constexpr NA2dDim value = {8, 16}; +//}; +// +//// 3D problems +// template <> +// struct MapTileSizeToNd<32, 3> { +// static constexpr NA3dDim value = {2, 2, 8}; +//}; +// +// template <> +// struct MapTileSizeToNd<64, 3> { +// static constexpr NA3dDim value = {4, 4, 4}; +//}; +// +// template <> +// struct MapTileSizeToNd<128, 3> { +// static constexpr NA3dDim value = {4, 4, 8}; +//}; + +// template +// struct Divmod; +// +// template <> +// struct Divmod<1> { +// using Dim = typename GetDim<1>::type; +// +// cutlass::FastDivmod x; +// +// CUTLASS_HOST_DEVICE constexpr Divmod(Dim divisor): x((int)divisor.x) {} +// +// CUTLASS_HOST_DEVICE void operator()(Dim "ient, Dim &remainder, Dim +// dividend) const { +// x(quotient.x, remainder.x, dividend.x); +// } +// }; +// +// template <> +// struct Divmod<2> { +// using Dim = typename GetDim<2>::type; +// +// cutlass::FastDivmod x; +// cutlass::FastDivmod y; +// +// CUTLASS_HOST_DEVICE constexpr Divmod(Dim divisor): x((int)divisor.x), +// y((int)divisor.y) {} +// +// CUTLASS_HOST_DEVICE void operator()(Dim "ient, Dim &remainder, Dim +// dividend) const { +// x(quotient.x, remainder.x, dividend.x); +// y(quotient.y, remainder.y, dividend.y); +// } +// }; +// +// template <> +// struct Divmod<3> { +// using Dim = typename GetDim<3>::type; +// +// cutlass::FastDivmod x; +// cutlass::FastDivmod y; +// cutlass::FastDivmod z; +// +// CUTLASS_HOST_DEVICE constexpr Divmod(Dim divisor): x((int)divisor.x), +// y((int)divisor.y) {} +// +// CUTLASS_HOST_DEVICE void operator()(Dim "ient, Dim &remainder, Dim +// dividend) const { +// x(quotient.x, remainder.x, dividend.x); +// y(quotient.y, remainder.y, dividend.y); +// } +// }; + +// template +// struct CoordMapper; +// +// template <> +// struct CoordMapper <1> { +// using Dim = typename GetDim<1>::type; +// +// CUTLASS_HOST_DEVICE CoordMapper(Dim block_shape) {} +// +// CUTLASS_HOST_DEVICE Dim operator()(int32_t block_index) { +// return Dim(block_index); +// } +//}; +// +// template <> +// struct CoordMapper <2> { +// using Dim = typename GetDim<2>::type; +// +// cutlass::FastDivmod divmod_y; +// +// CUTLASS_HOST_DEVICE CoordMapper(Dim block_shape): +// divmod_y((int)block_shape.y) {} +// +// CUTLASS_HOST_DEVICE Dim operator()(int32_t block_index) { +// int index_x, index_y; +// divmod_y(index_x, index_y, block_index); +// return Dim(index_x, index_y); +// } +//}; +// +// template <> +// struct CoordMapper <3> { +// using Dim = typename GetDim<3>::type; +// +// cutlass::FastDivmod divmod_y; +// cutlass::FastDivmod divmod_z; +// +// CUTLASS_HOST_DEVICE CoordMapper(Dim block_shape): +// divmod_y((int)block_shape.y), divmod_z((int)block_shape.z) {} +// +// CUTLASS_HOST_DEVICE Dim operator()(int32_t block_index) { +// int index_x, index_y, index_z, index_xy; +// divmod_z(index_xy, index_z, block_index); +// divmod_z(index_x, index_y, index_xy); +// return Dim(index_x, index_y, index_z); +// } +//}; + +template +CUTLASS_HOST_DEVICE Dim +map_index_to_coord(int32_t block_index, Dim block_shape); + +template <> +CUTLASS_HOST_DEVICE NA1dDim +map_index_to_coord(int32_t block_index, NA1dDim block_shape) { + // 1d is just linear indexing. + return NA1dDim(block_index); +}; + +template <> +CUTLASS_HOST_DEVICE NA2dDim +map_index_to_coord(int32_t block_index, NA2dDim block_shape) { + auto index_x = block_index / block_shape.y; + return NA2dDim(index_x, block_index - index_x * block_shape.y); +} + +template <> +CUTLASS_HOST_DEVICE NA3dDim +map_index_to_coord(int32_t block_index, NA3dDim block_shape) { + auto index_z = block_index % block_shape.z; + auto index_xy = block_index / block_shape.z; + auto index_y = index_xy % block_shape.y; + auto index_x = index_xy / block_shape.y; + return NA3dDim(index_x, index_y, index_z); +} + +template +CUTLASS_HOST_DEVICE int32_t map_coord_to_index(Dim coord, Dim block_shape); + +template <> +CUTLASS_HOST_DEVICE int32_t +map_coord_to_index(NA1dDim coord, NA1dDim block_shape) { + // 1d is just linear indexing. + return coord.x; +} + +template <> +CUTLASS_HOST_DEVICE int32_t +map_coord_to_index(NA2dDim coord, NA2dDim block_shape) { + return coord.x * block_shape.y + coord.y; +} + +template <> +CUTLASS_HOST_DEVICE int32_t +map_coord_to_index(NA3dDim coord, NA3dDim block_shape) { + return coord.x * block_shape.y * block_shape.z + coord.y * block_shape.z + + coord.z; +} + +template +CUTLASS_HOST_DEVICE Dim fast_min(Dim lhs, Dim rhs); + +template <> +CUTLASS_HOST_DEVICE NA1dDim fast_min(NA1dDim lhs, NA1dDim rhs) { + return NA1dDim(cutlass::fast_min(lhs.x, rhs.x)); +} + +template <> +CUTLASS_HOST_DEVICE NA2dDim fast_min(NA2dDim lhs, NA2dDim rhs) { + return NA2dDim( + cutlass::fast_min(lhs.x, rhs.x), cutlass::fast_min(lhs.y, rhs.y)); +} + +template <> +CUTLASS_HOST_DEVICE NA3dDim fast_min(NA3dDim lhs, NA3dDim rhs) { + return NA3dDim( + cutlass::fast_min(lhs.x, rhs.x), + cutlass::fast_min(lhs.y, rhs.y), + cutlass::fast_min(lhs.z, rhs.z)); +} + +template +CUTLASS_HOST_DEVICE Dim fast_max(Dim lhs, Dim rhs); + +template <> +CUTLASS_HOST_DEVICE NA1dDim fast_max(NA1dDim lhs, NA1dDim rhs) { + return NA1dDim(cutlass::fast_max(lhs.x, rhs.x)); +} + +template <> +CUTLASS_HOST_DEVICE NA2dDim fast_max(NA2dDim lhs, NA2dDim rhs) { + return NA2dDim( + cutlass::fast_max(lhs.x, rhs.x), cutlass::fast_max(lhs.y, rhs.y)); +} + +template <> +CUTLASS_HOST_DEVICE NA3dDim fast_max(NA3dDim lhs, NA3dDim rhs) { + return NA3dDim( + cutlass::fast_max(lhs.x, rhs.x), + cutlass::fast_max(lhs.y, rhs.y), + cutlass::fast_max(lhs.z, rhs.z)); +} + +template +CUTLASS_HOST_DEVICE Dim ceil_div_dim(Dim lhs, Dim rhs) {} + +template <> +CUTLASS_HOST_DEVICE NA1dDim ceil_div_dim(NA1dDim lhs, NA1dDim rhs) { + return NA1dDim(gemm_kernel_utils::ceil_div(lhs.x, rhs.x)); +} + +template <> +CUTLASS_HOST_DEVICE NA2dDim ceil_div_dim(NA2dDim lhs, NA2dDim rhs) { + return NA2dDim( + gemm_kernel_utils::ceil_div(lhs.x, rhs.x), + gemm_kernel_utils::ceil_div(lhs.y, rhs.y)); +} + +template <> +CUTLASS_HOST_DEVICE NA3dDim ceil_div_dim(NA3dDim lhs, NA3dDim rhs) { + return NA3dDim( + gemm_kernel_utils::ceil_div(lhs.x, rhs.x), + gemm_kernel_utils::ceil_div(lhs.y, rhs.y), + gemm_kernel_utils::ceil_div(lhs.z, rhs.z)); +} + +template +CUTLASS_HOST_DEVICE Dim div_dim(Dim lhs, Dim rhs) {} + +template <> +CUTLASS_HOST_DEVICE NA1dDim div_dim(NA1dDim lhs, NA1dDim rhs) { + return NA1dDim(lhs.x / rhs.x); +} + +template <> +CUTLASS_HOST_DEVICE NA2dDim div_dim(NA2dDim lhs, NA2dDim rhs) { + return NA2dDim(lhs.x / rhs.x, lhs.y / rhs.y); +} + +template <> +CUTLASS_HOST_DEVICE NA3dDim div_dim(NA3dDim lhs, NA3dDim rhs) { + return NA3dDim(lhs.x / rhs.x, lhs.y / rhs.y, lhs.z / rhs.z); +} + +template +CUTLASS_HOST_DEVICE Dim compute_stride(Dim input_shape, int64_t last_stride); + +template <> +CUTLASS_HOST_DEVICE NA1dDim +compute_stride(NA1dDim input_shape, int64_t last_stride) { + return NA1dDim(last_stride); +} + +template <> +CUTLASS_HOST_DEVICE NA2dDim +compute_stride(NA2dDim input_shape, int64_t last_stride) { + return NA2dDim(input_shape.y * last_stride, last_stride); +} + +template <> +CUTLASS_HOST_DEVICE NA3dDim +compute_stride(NA3dDim input_shape, int64_t last_stride) { + return NA3dDim( + input_shape.y * input_shape.z * last_stride, + input_shape.z * last_stride, + last_stride); +} + +template +CUTLASS_HOST_DEVICE bool is_coord_within_upper_bound( + Dim coord, + Dim upper_bound); + +template <> +CUTLASS_HOST_DEVICE bool is_coord_within_upper_bound( + NA1dDim coord, + NA1dDim upper_bound) { + return coord.x < upper_bound.x; +} + +template <> +CUTLASS_HOST_DEVICE bool is_coord_within_upper_bound( + NA2dDim coord, + NA2dDim upper_bound) { + return coord.x < upper_bound.x && coord.y < upper_bound.y; +} + +template <> +CUTLASS_HOST_DEVICE bool is_coord_within_upper_bound( + NA3dDim coord, + NA3dDim upper_bound) { + return coord.x < upper_bound.x && coord.y < upper_bound.y && + coord.z < upper_bound.z; +} + +template +CUTLASS_HOST_DEVICE bool is_coord_within_bounds( + Dim coord, + Dim lower_bound, + Dim upper_bound); + +template <> +CUTLASS_HOST_DEVICE bool is_coord_within_bounds( + NA1dDim coord, + NA1dDim lower_bound, + NA1dDim upper_bound) { + return coord.x >= lower_bound.x && coord.x < upper_bound.x; +} + +template <> +CUTLASS_HOST_DEVICE bool is_coord_within_bounds( + NA2dDim coord, + NA2dDim lower_bound, + NA2dDim upper_bound) { + return coord.x >= lower_bound.x && coord.x < upper_bound.x && + coord.y >= lower_bound.y && coord.y < upper_bound.y; +} + +template <> +CUTLASS_HOST_DEVICE bool is_coord_within_bounds( + NA3dDim coord, + NA3dDim lower_bound, + NA3dDim upper_bound) { + return coord.x >= lower_bound.x && coord.x < upper_bound.x && + coord.y >= lower_bound.y && coord.y < upper_bound.y && + coord.z >= lower_bound.z && coord.z < upper_bound.z; +} + +template +CUTLASS_HOST_DEVICE bool is_coord_within_bounds_nn(Dim coord, Dim upper_bound); + +template <> +CUTLASS_HOST_DEVICE bool is_coord_within_bounds_nn( + NA1dDim coord, + NA1dDim upper_bound) { + return coord.x >= 0 && coord.x < upper_bound.x; +} + +template <> +CUTLASS_HOST_DEVICE bool is_coord_within_bounds_nn( + NA2dDim coord, + NA2dDim upper_bound) { + return coord.x >= 0 && coord.y >= 0 && coord.x < upper_bound.x && + coord.y < upper_bound.y; +} + +template <> +CUTLASS_HOST_DEVICE bool is_coord_within_bounds_nn( + NA3dDim coord, + NA3dDim upper_bound) { + return coord.x >= 0 && coord.y >= 0 && coord.z >= 0 && + coord.x < upper_bound.x && coord.y < upper_bound.y && + coord.z < upper_bound.z; +} + +template +CUTLASS_HOST_DEVICE Dim align_dim_down(Dim tile, Dim upper_bound); +template +CUTLASS_HOST_DEVICE Dim align_dim_up(Dim tile, Dim upper_bound); + +template <> +CUTLASS_HOST_DEVICE NA1dDim align_dim_down(NA1dDim tile, NA1dDim upper_bound) { + return NA1dDim((tile.x / upper_bound.x) * upper_bound.x); +} +template <> +CUTLASS_HOST_DEVICE NA1dDim align_dim_up(NA1dDim tile, NA1dDim upper_bound) { + return NA1dDim( + ((tile.x + upper_bound.x - 1) / upper_bound.x) * upper_bound.x); +} + +template <> +CUTLASS_HOST_DEVICE NA2dDim align_dim_down(NA2dDim tile, NA2dDim upper_bound) { + return NA2dDim( + (tile.x / upper_bound.x) * upper_bound.x, + (tile.y / upper_bound.y) * upper_bound.y); +} +template <> +CUTLASS_HOST_DEVICE NA2dDim align_dim_up(NA2dDim tile, NA2dDim upper_bound) { + return NA2dDim( + ((tile.x + upper_bound.x - 1) / upper_bound.x) * upper_bound.x, + ((tile.y + upper_bound.y - 1) / upper_bound.y) * upper_bound.y); +} + +template <> +CUTLASS_HOST_DEVICE NA3dDim align_dim_down(NA3dDim tile, NA3dDim upper_bound) { + return NA3dDim( + (tile.x / upper_bound.x) * upper_bound.x, + (tile.y / upper_bound.y) * upper_bound.y, + (tile.z / upper_bound.z) * upper_bound.z); +} +template <> +CUTLASS_HOST_DEVICE NA3dDim align_dim_up(NA3dDim tile, NA3dDim upper_bound) { + return NA3dDim( + ((tile.x + upper_bound.x - 1) / upper_bound.x) * upper_bound.x, + ((tile.y + upper_bound.y - 1) / upper_bound.y) * upper_bound.y, + ((tile.z + upper_bound.z - 1) / upper_bound.z) * upper_bound.z); +} + +template +CUTLASS_HOST_DEVICE bool is_coord_less_than_or_equal_to( + Dim coord, + Dim upper_bound); + +template <> +CUTLASS_HOST_DEVICE bool is_coord_less_than_or_equal_to( + NA1dDim coord, + NA1dDim upper_bound) { + return coord.x <= upper_bound.x; +} + +template <> +CUTLASS_HOST_DEVICE bool is_coord_less_than_or_equal_to( + NA2dDim coord, + NA2dDim upper_bound) { + return coord.x <= upper_bound.x && coord.y <= upper_bound.y; +} + +template <> +CUTLASS_HOST_DEVICE bool is_coord_less_than_or_equal_to( + NA3dDim coord, + NA3dDim upper_bound) { + return coord.x <= upper_bound.x && coord.y <= upper_bound.y && + coord.z <= upper_bound.z; +} + +template +CUTLASS_HOST_DEVICE bool is_coord_greater_than_or_equal_to( + Dim coord, + Dim upper_bound); + +template <> +CUTLASS_HOST_DEVICE bool is_coord_greater_than_or_equal_to( + NA1dDim coord, + NA1dDim upper_bound) { + return coord.x >= upper_bound.x; +} + +template <> +CUTLASS_HOST_DEVICE bool is_coord_greater_than_or_equal_to( + NA2dDim coord, + NA2dDim upper_bound) { + return coord.x >= upper_bound.x && coord.y >= upper_bound.y; +} + +template <> +CUTLASS_HOST_DEVICE bool is_coord_greater_than_or_equal_to( + NA3dDim coord, + NA3dDim upper_bound) { + return coord.x >= upper_bound.x && coord.y >= upper_bound.y && + coord.z >= upper_bound.z; +} + +// Device only + +template +CUTLASS_DEVICE void maybe_mask_qk_tiles( + Dim& q_shape, + Dim orig_shape, + Dim dilation, + Dim dilation_idx); + +template <> +CUTLASS_DEVICE void maybe_mask_qk_tiles( + NA1dDim& q_shape, + NA1dDim orig_shape, + NA1dDim dilation, + NA1dDim dilation_idx) { + auto num_tiles_with_fewer_tokens = (q_shape * dilation) - orig_shape; + // TODO: assert num_tiles_with_fewer_tokens >= 0 + if (dilation_idx.x >= dilation.x - num_tiles_with_fewer_tokens.x) { + q_shape.x = q_shape.x - 1; + } +} + +template <> +CUTLASS_DEVICE void maybe_mask_qk_tiles( + NA2dDim& q_shape, + NA2dDim orig_shape, + NA2dDim dilation, + NA2dDim dilation_idx) { + auto num_tiles_with_fewer_tokens = (q_shape * dilation) - orig_shape; + // TODO: assert num_tiles_with_fewer_tokens >= 0 + if (dilation_idx.x >= dilation.x - num_tiles_with_fewer_tokens.x) { + q_shape.x = q_shape.x - 1; + } + if (dilation_idx.y >= dilation.y - num_tiles_with_fewer_tokens.y) { + q_shape.y = q_shape.y - 1; + } +} + +template <> +CUTLASS_DEVICE void maybe_mask_qk_tiles( + NA3dDim& q_shape, + NA3dDim orig_shape, + NA3dDim dilation, + NA3dDim dilation_idx) { + auto num_tiles_with_fewer_tokens = (q_shape * dilation) - orig_shape; + // TODO: assert num_tiles_with_fewer_tokens >= 0 + if (dilation_idx.x >= dilation.x - num_tiles_with_fewer_tokens.x) { + q_shape.x = q_shape.x - 1; + } + if (dilation_idx.y >= dilation.y - num_tiles_with_fewer_tokens.y) { + q_shape.y = q_shape.y - 1; + } + if (dilation_idx.z >= dilation.z - num_tiles_with_fewer_tokens.z) { + q_shape.z = q_shape.z - 1; + } +} + +template +struct CausalMask; + +template +struct CausalMask { + static constexpr int kRank = 1; + static constexpr bool Dim0 = IsDim0Causal; + + static constexpr bool AnyCausalDims = IsDim0Causal; +}; + +template +struct CausalMask { + static constexpr int kRank = 2; + static constexpr bool Dim0 = IsDim0Causal; + static constexpr bool Dim1 = IsDim1Causal; + + static constexpr bool AnyCausalDims = IsDim0Causal || IsDim1Causal; +}; + +template +struct CausalMask { + static constexpr int kRank = 3; + static constexpr bool Dim0 = IsDim0Causal; + static constexpr bool Dim1 = IsDim1Causal; + static constexpr bool Dim2 = IsDim2Causal; + + static constexpr bool AnyCausalDims = + IsDim0Causal || IsDim1Causal || IsDim2Causal; +}; + +// No causal masking +template +struct GetNonCausalMask; + +template <> +struct GetNonCausalMask<1> { + using type = CausalMask; +}; + +template <> +struct GetNonCausalMask<2> { + using type = CausalMask; +}; + +template <> +struct GetNonCausalMask<3> { + using type = CausalMask; +}; + +// NA Mask +template +struct NeighborhoodAttentionMaskBase { + static constexpr bool DoCausalMasking = DoCausalMasking_; +}; + +// Standard NA + +// TODO: duplicated for fully block sparse check +CUTLASS_DEVICE +int neighborhood_atten_get_window_start( + int index, + int window_radius_left, + int window_radius_right, + int stride, + int spatial_extent) { + // Stride group leader is the center-most query in the stride group. + // If stride is even, choose the right hand side center query. + int stride_group_leader_idx = cutlass::fast_min( + ((index / stride) * stride) + (stride / 2), spatial_extent - 1); + return cutlass::fast_max(stride_group_leader_idx - window_radius_left, 0) + + ((stride_group_leader_idx + window_radius_right >= spatial_extent) * + (spatial_extent - window_radius_right - stride_group_leader_idx - 1)); +} + +template <> +struct NeighborhoodAttentionMaskBase { + static constexpr bool DoCausalMasking = false; + + int32_t window_size; + int32_t window_radius_left; + int32_t window_radius_right; + int32_t stride; + int32_t stride_group_center_offset; + int32_t spatial_extent; + + CUTLASS_DEVICE NeighborhoodAttentionMaskBase( + int32_t kernel_size, + int32_t stride_, + int32_t spatial_extent_) + : window_size(kernel_size), + stride(stride_), + window_radius_left(kernel_size / 2), + window_radius_right((kernel_size / 2) + ((kernel_size % 2) - 1)), + stride_group_center_offset(stride_ - (stride_ / 2) - 1), + spatial_extent(spatial_extent_) {} + + CUTLASS_DEVICE int32_t get_window_start(int32_t index) const { + // Stride group leader is the center-most query in the stride group. + // If stride is even, choose the right hand side center query. + int32_t stride_group_leader_idx = cutlass::fast_min( + ((index / stride) * stride) + (stride / 2), spatial_extent - 1); + return cutlass::fast_max(stride_group_leader_idx - window_radius_left, 0) + + ((stride_group_leader_idx + window_radius_right >= spatial_extent) * + (spatial_extent - window_radius_right - stride_group_leader_idx - 1)); + } + + CUTLASS_DEVICE int32_t get_window_end(int32_t index, int32_t start) const { + return start + window_size; + } + + CUTLASS_DEVICE int32_t get_window_end_(int32_t index) const { + return get_window_end(index, get_window_start(index)); + } + + CUTLASS_DEVICE int32_t get_backward_window_start(int32_t index) const { + int32_t window_start_pre_trunc = + (index >= (window_radius_left + window_radius_right + 1)) * + (index - window_radius_right + stride_group_center_offset); + return (window_start_pre_trunc / stride) * stride; + } + + CUTLASS_DEVICE int32_t get_backward_window_end(int32_t index) const { + int32_t extent_check = + (index >= + (spatial_extent - window_radius_left - window_radius_right - 1)); + int32_t window_end = + ((index + window_radius_left + stride_group_center_offset + 1) / + stride) * + stride; + return (extent_check * spatial_extent) + ((1 - extent_check) * window_end); + } +}; + +// Causal NA +template <> +struct NeighborhoodAttentionMaskBase { + static constexpr bool DoCausalMasking = true; + + int32_t window_size; + int32_t stride; + int32_t spatial_extent; + + // NOTE: Causal only supports stride=1 + CUTLASS_DEVICE NeighborhoodAttentionMaskBase( + int32_t kernel_size, + int32_t stride_, + int32_t spatial_extent_) + : window_size(kernel_size), + stride(stride_), + spatial_extent(spatial_extent_) {} + + CUTLASS_DEVICE int32_t get_window_start(int32_t index) const { + // Stride group leader is the last (right-most) query in the stride group. + int32_t stride_group_leader_idx = cutlass::fast_min( + (index / stride) * stride + stride - 1, spatial_extent - 1); + return cutlass::fast_max(stride_group_leader_idx - window_size + 1, 0); + } + + CUTLASS_DEVICE int32_t get_window_end(int32_t index, int32_t start) const { + return cutlass::fast_min(index + 1, spatial_extent); + } + + CUTLASS_DEVICE int32_t get_window_end_(int32_t index) const { + return cutlass::fast_min(index + 1, spatial_extent); + } + + CUTLASS_DEVICE int32_t get_backward_window_start(int32_t index) const { + return index; + } + + CUTLASS_DEVICE int32_t get_backward_window_end(int32_t index) const { + // Window end is always the last query that attends to this key/value + 1 + // In the strided case, it will always be the last query in a stride + // group + 1, which means it will always be a multiple of stride. + // In the case of incomplete stridegroups (input size % stride != 0), + // we must correct backward window end by incrementing first non-attending + // stridegroup by 1. + auto first_non_attending_stride_group_idx = (index + window_size) / stride + + (index + window_size >= spatial_extent); + return cutlass::fast_min( + first_non_attending_stride_group_idx * stride, spatial_extent); + } +}; + +// TODO: figure out an alternative to partial specializations for +// literally every dim. I'm pretty sure there's a better way of handling +// it with less repetition. It might require refactoring this whole +// header file, particularly how NA.dDim structs are defined. + +CUTLASS_DEVICE int32_t align_down(int32_t tile, int32_t upper_bound) { + return (tile / upper_bound) * upper_bound; +} + +template +struct NeighborhoodAttentionMask; + +template +struct NeighborhoodAttentionMask { + static_assert(NADim >= 1 && NADim < 4); + static_assert(NADim == CausalMask_::kRank); + + using CausalMask = CausalMask_; +}; + +template +struct NeighborhoodAttentionMask<1, CausalMask_> { + static constexpr int kRank = 1; + + using MaskType0 = NeighborhoodAttentionMaskBase; + using Dim = typename natten::cuda::fna::GetDim::type; + + MaskType0 mask_0; + + CUTLASS_DEVICE NeighborhoodAttentionMask( + Dim kernel_size, + Dim stride, + Dim spatial_extent_) + : mask_0(kernel_size.x, stride.x, spatial_extent_.x) {} + + CUTLASS_DEVICE Dim get_window_start(Dim index) const { + return Dim(mask_0.get_window_start(index.x)); + } + + CUTLASS_DEVICE Dim get_window_end(Dim index, Dim start) const { + return Dim(mask_0.get_window_end(index.x, start.x)); + } + + CUTLASS_DEVICE Dim get_window_end_(Dim index) const { + return Dim(mask_0.get_window_end_(index.x)); + } + + CUTLASS_DEVICE Dim get_last_key_block_for_query_block( + Dim index, + Dim key_tile_shape, + Dim last_possible_key_tile) const { + return Dim(cutlass::fast_min( + align_down(mask_0.get_window_end_(index.x) - 1, key_tile_shape.x), + last_possible_key_tile.x)); + } + + CUTLASS_DEVICE Dim get_backward_window_start(Dim index) const { + return Dim(mask_0.get_backward_window_start(index.x)); + } + + CUTLASS_DEVICE Dim get_backward_window_end(Dim index) const { + return Dim(mask_0.get_backward_window_end(index.x)); + } +}; + +template +struct NeighborhoodAttentionMask<2, CausalMask_> { + static constexpr int kRank = 2; + + using MaskType0 = NeighborhoodAttentionMaskBase; + using MaskType1 = NeighborhoodAttentionMaskBase; + using Dim = typename natten::cuda::fna::GetDim::type; + + MaskType0 mask_0; + MaskType1 mask_1; + + CUTLASS_DEVICE NeighborhoodAttentionMask( + Dim kernel_size, + Dim stride, + Dim spatial_extent_) + : mask_0(kernel_size.x, stride.x, spatial_extent_.x), + mask_1(kernel_size.y, stride.y, spatial_extent_.y) {} + + CUTLASS_DEVICE Dim get_window_start(Dim index) const { + return Dim( + mask_0.get_window_start(index.x), mask_1.get_window_start(index.y)); + } + + CUTLASS_DEVICE Dim get_window_end(Dim index, Dim start) const { + return Dim( + mask_0.get_window_end(index.x, start.x), + mask_1.get_window_end(index.y, start.y)); + } + + CUTLASS_DEVICE Dim get_window_end_(Dim index) const { + return Dim( + mask_0.get_window_end_(index.x), mask_1.get_window_end_(index.y)); + } + + CUTLASS_DEVICE Dim get_last_key_block_for_query_block( + Dim index, + Dim key_tile_shape, + Dim last_possible_key_tile) const { + return Dim( + cutlass::fast_min( + align_down(mask_0.get_window_end_(index.x) - 1, key_tile_shape.x), + last_possible_key_tile.x), + cutlass::fast_min( + align_down(mask_1.get_window_end_(index.y) - 1, key_tile_shape.y), + last_possible_key_tile.y)); + } + + CUTLASS_DEVICE Dim get_backward_window_start(Dim index) const { + return Dim( + mask_0.get_backward_window_start(index.x), + mask_1.get_backward_window_start(index.y)); + } + + CUTLASS_DEVICE Dim get_backward_window_end(Dim index) const { + return Dim( + mask_0.get_backward_window_end(index.x), + mask_1.get_backward_window_end(index.y)); + } +}; + +template +struct NeighborhoodAttentionMask<3, CausalMask_> { + static constexpr int kRank = 3; + + using MaskType0 = NeighborhoodAttentionMaskBase; + using MaskType1 = NeighborhoodAttentionMaskBase; + using MaskType2 = NeighborhoodAttentionMaskBase; + using Dim = typename natten::cuda::fna::GetDim::type; + + MaskType0 mask_0; + MaskType1 mask_1; + MaskType2 mask_2; + + CUTLASS_DEVICE NeighborhoodAttentionMask( + Dim kernel_size, + Dim stride, + Dim spatial_extent_) + : mask_0(kernel_size.x, stride.x, spatial_extent_.x), + mask_1(kernel_size.y, stride.y, spatial_extent_.y), + mask_2(kernel_size.z, stride.z, spatial_extent_.z) {} + + CUTLASS_DEVICE Dim get_window_start(Dim index) const { + return Dim( + mask_0.get_window_start(index.x), + mask_1.get_window_start(index.y), + mask_2.get_window_start(index.z)); + } + + CUTLASS_DEVICE Dim get_window_end(Dim index, Dim start) const { + return Dim( + mask_0.get_window_end(index.x, start.x), + mask_1.get_window_end(index.y, start.y), + mask_2.get_window_end(index.z, start.z)); + } + + CUTLASS_DEVICE Dim get_window_end_(Dim index) const { + return Dim( + mask_0.get_window_end_(index.x), + mask_1.get_window_end_(index.y), + mask_2.get_window_end_(index.z)); + } + + CUTLASS_DEVICE Dim get_last_key_block_for_query_block( + Dim index, + Dim key_tile_shape, + Dim last_possible_key_tile) const { + return Dim( + cutlass::fast_min( + align_down(mask_0.get_window_end_(index.x) - 1, key_tile_shape.x), + last_possible_key_tile.x), + cutlass::fast_min( + align_down(mask_1.get_window_end_(index.y) - 1, key_tile_shape.y), + last_possible_key_tile.y), + cutlass::fast_min( + align_down(mask_2.get_window_end_(index.z) - 1, key_tile_shape.z), + last_possible_key_tile.z)); + } + + CUTLASS_DEVICE Dim get_backward_window_start(Dim index) const { + return Dim( + mask_0.get_backward_window_start(index.x), + mask_1.get_backward_window_start(index.y), + mask_2.get_backward_window_start(index.z)); + } + + CUTLASS_DEVICE Dim get_backward_window_end(Dim index) const { + return Dim( + mask_0.get_backward_window_end(index.x), + mask_1.get_backward_window_end(index.y), + mask_2.get_backward_window_end(index.z)); + } +}; + +template +CUTLASS_DEVICE void increment_tile( + Dim& tile, + const Dim& inc_value, + const Dim& upper_bound); + +template <> +CUTLASS_DEVICE void increment_tile( + NA1dDim& tile, + const NA1dDim& inc_value, + const NA1dDim& upper_bound) { + tile.x += inc_value.x; +} + +template <> +CUTLASS_DEVICE void increment_tile( + NA2dDim& tile, + const NA2dDim& inc_value, + const NA2dDim& upper_bound) { + tile.y = tile.y + inc_value.y; + tile.x = tile.x + (tile.y >= upper_bound.y) * inc_value.x; + tile.y = (tile.y < upper_bound.y) * tile.y; +} + +template <> +CUTLASS_DEVICE void increment_tile( + NA3dDim& tile, + const NA3dDim& inc_value, + const NA3dDim& upper_bound) { + tile.z = tile.z + inc_value.z; + if (tile.z >= upper_bound.z) { + tile.z = 0; + tile.y += inc_value.y; + } + if (tile.y >= upper_bound.y) { + tile.y = 0; + tile.x += inc_value.x; + } +} + +template +CUTLASS_DEVICE void increment_tile( + Dim& tile, + const Dim& inc_value, + const Dim& lower_bound, + const Dim& upper_bound); + +template <> +CUTLASS_DEVICE void increment_tile( + NA1dDim& tile, + const NA1dDim& inc_value, + const NA1dDim& lower_bound, + const NA1dDim& upper_bound) { + tile.x += inc_value.x; +} + +template <> +CUTLASS_DEVICE void increment_tile( + NA2dDim& tile, + const NA2dDim& inc_value, + const NA2dDim& lower_bound, + const NA2dDim& upper_bound) { + tile.y = tile.y + inc_value.y; + if (tile.y >= upper_bound.y) { + tile.y = lower_bound.y; + tile.x += inc_value.x; + } +} + +template <> +CUTLASS_DEVICE void increment_tile( + NA3dDim& tile, + const NA3dDim& inc_value, + const NA3dDim& lower_bound, + const NA3dDim& upper_bound) { + tile.z = tile.z + inc_value.z; + if (tile.z >= upper_bound.z) { + tile.z = lower_bound.z; + tile.y += inc_value.y; + } + if (tile.y >= upper_bound.y) { + tile.y = lower_bound.y; + tile.x += inc_value.x; + } +} + +template +CUTLASS_DEVICE Dim increment_tile_return( + const Dim& tile, + const Dim& inc_value, + const Dim& lower_bound, + const Dim& upper_bound); +template +CUTLASS_DEVICE Dim increment_tile_return( + const Dim& tile, + const Dim& inc_value, + const Dim& upper_bound); + +template <> +CUTLASS_DEVICE NA1dDim increment_tile_return( + const NA1dDim& tile, + const NA1dDim& inc_value, + const NA1dDim& lower_bound, + const NA1dDim& upper_bound) { + return NA1dDim(tile.x + inc_value.x); +} + +template <> +CUTLASS_DEVICE NA1dDim increment_tile_return( + const NA1dDim& tile, + const NA1dDim& inc_value, + const NA1dDim& upper_bound) { + return NA1dDim(tile.x + inc_value.x); +} + +template <> +CUTLASS_DEVICE NA2dDim increment_tile_return( + const NA2dDim& tile, + const NA2dDim& inc_value, + const NA2dDim& lower_bound, + const NA2dDim& upper_bound) { + auto x = tile.x; + auto y = tile.y + inc_value.y; + if (y >= upper_bound.y) { + y = lower_bound.y; + x += inc_value.x; + } + return NA2dDim(x, y); +} + +template <> +CUTLASS_DEVICE NA2dDim increment_tile_return( + const NA2dDim& tile, + const NA2dDim& inc_value, + const NA2dDim& upper_bound) { + auto x = tile.x; + auto y = tile.y + inc_value.y; + if (y >= upper_bound.y) { + y = 0; + x += inc_value.x; + } + return NA2dDim(x, y); +} + +template <> +CUTLASS_DEVICE NA3dDim increment_tile_return( + const NA3dDim& tile, + const NA3dDim& inc_value, + const NA3dDim& lower_bound, + const NA3dDim& upper_bound) { + auto x = tile.x; + auto y = tile.y; + auto z = tile.z + inc_value.z; + if (z >= upper_bound.z) { + z = lower_bound.z; + y += inc_value.y; + } + if (y >= upper_bound.y) { + y = lower_bound.y; + x += inc_value.x; + } + return NA3dDim(x, y, z); +} + +template <> +CUTLASS_DEVICE NA3dDim increment_tile_return( + const NA3dDim& tile, + const NA3dDim& inc_value, + const NA3dDim& upper_bound) { + auto x = tile.x; + auto y = tile.y; + auto z = tile.z + inc_value.z; + if (z >= upper_bound.z) { + z = 0; + y += inc_value.y; + } + if (y >= upper_bound.y) { + y = 0; + x += inc_value.x; + } + return NA3dDim(x, y, z); +} + +template +CUTLASS_DEVICE bool is_tile_last(Dim tile, Dim inc_value, Dim upper_bound); + +template <> +CUTLASS_DEVICE bool is_tile_last( + NA1dDim tile, + NA1dDim inc_value, + NA1dDim upper_bound) { + return tile.x + inc_value.x >= upper_bound.x; +} + +template <> +CUTLASS_DEVICE bool is_tile_last( + NA2dDim tile, + NA2dDim inc_value, + NA2dDim upper_bound) { + return tile.y + inc_value.y >= upper_bound.y && + tile.x + inc_value.x >= upper_bound.x; +} + +template <> +CUTLASS_DEVICE bool is_tile_last( + NA3dDim tile, + NA3dDim inc_value, + NA3dDim upper_bound) { + return tile.z + inc_value.z >= upper_bound.z && + tile.y + inc_value.y >= upper_bound.y && + tile.x + inc_value.x >= upper_bound.x; +} + +template +CUTLASS_DEVICE Dim count_tiles(Dim first, Dim last, Dim tile_shape); + +template <> +CUTLASS_DEVICE NA1dDim +count_tiles(NA1dDim first, NA1dDim last, NA1dDim tile_shape) { + return NA1dDim((last.x / tile_shape.x) - (first.x / tile_shape.x) + 1); +} + +template <> +CUTLASS_DEVICE NA2dDim +count_tiles(NA2dDim first, NA2dDim last, NA2dDim tile_shape) { + return NA2dDim( + ((last.x / tile_shape.x) - (first.x / tile_shape.x) + 1), + ((last.y / tile_shape.y) - (first.y / tile_shape.y) + 1)); +} + +template <> +CUTLASS_DEVICE NA3dDim +count_tiles(NA3dDim first, NA3dDim last, NA3dDim tile_shape) { + return NA3dDim( + ((last.x / tile_shape.x) - (first.x / tile_shape.x) + 1), + ((last.y / tile_shape.y) - (first.y / tile_shape.y) + 1), + ((last.z / tile_shape.z) - (first.z / tile_shape.z) + 1)); +} + +//// Block sparse fast path + +template +CUTLASS_DEVICE bool evenly_divides(Dim a, Dim b); + +template <> +CUTLASS_DEVICE bool evenly_divides(NA1dDim a, NA1dDim b) { + return a.x % b.x == 0; +} + +template <> +CUTLASS_DEVICE bool evenly_divides(NA2dDim a, NA2dDim b) { + return a.x % b.x == 0 && a.y % b.y == 0; +} + +template <> +CUTLASS_DEVICE bool evenly_divides(NA3dDim a, NA3dDim b) { + return a.x % b.x == 0 && a.y % b.y == 0; +} + +// NOTE: this is for bi-directional ONLY! +CUTLASS_DEVICE +bool is_fully_block_sparse_int( + int input_size, + int window_size, + int stride, + int q_tile_size, + int kv_tile_size) { + int num_stride_groups = gemm_kernel_utils::ceil_div(input_size, stride); + int window_left = window_size / 2; + int window_right = (window_size / 2) + ((window_size % 2) - 1); + int stride_group_center = stride / 2; + + int last_stride_group_center = + min((input_size / stride) * stride + stride_group_center, input_size - 1); + int last_stride_group_window_start = + neighborhood_atten_get_window_start /**/ ( // BI-DIRECTIONAL + // MASK + last_stride_group_center, + window_left, + window_right, + stride, + input_size); + + return input_size == window_size || + (stride % q_tile_size == 0 && window_size % kv_tile_size == 0 && + (input_size % stride == 0 || + last_stride_group_window_start % kv_tile_size == 0) && + (((window_left - stride_group_center) % kv_tile_size == 0) || + ( + // input_size % stride == 0 && + input_size % kv_tile_size == 0 && num_stride_groups == 2 && + stride + stride_group_center >= input_size - window_size))); +} + +// NOTE: this is for bi-directional ONLY! +// Causal masking along any dim will never be fully block sparse. +// In Hopper and Blackwell FNA, we check that within this function. +// Here it is assumed it's already checked. +template +CUTLASS_DEVICE bool fully_block_sparse( + Dim qkv_shape, + Dim window_size, + Dim stride, + Dim q_tile_shape, + Dim kv_tile_shape); + +template <> +CUTLASS_DEVICE bool fully_block_sparse( + NA1dDim qkv_shape, + NA1dDim window_size, + NA1dDim stride, + NA1dDim q_tile_shape, + NA1dDim kv_tile_shape) { + return is_fully_block_sparse_int( + qkv_shape.x, window_size.x, stride.x, q_tile_shape.x, kv_tile_shape.x); +} + +template <> +CUTLASS_DEVICE bool fully_block_sparse( + NA2dDim qkv_shape, + NA2dDim window_size, + NA2dDim stride, + NA2dDim q_tile_shape, + NA2dDim kv_tile_shape) { + return is_fully_block_sparse_int( + qkv_shape.x, + window_size.x, + stride.x, + q_tile_shape.x, + kv_tile_shape.x) && + is_fully_block_sparse_int( + qkv_shape.y, + window_size.y, + stride.y, + q_tile_shape.y, + kv_tile_shape.y); +} + +template <> +CUTLASS_DEVICE bool fully_block_sparse( + NA3dDim qkv_shape, + NA3dDim window_size, + NA3dDim stride, + NA3dDim q_tile_shape, + NA3dDim kv_tile_shape) { + return is_fully_block_sparse_int( + qkv_shape.x, + window_size.x, + stride.x, + q_tile_shape.x, + kv_tile_shape.x) && + is_fully_block_sparse_int( + qkv_shape.y, + window_size.y, + stride.y, + q_tile_shape.y, + kv_tile_shape.y) && + is_fully_block_sparse_int( + qkv_shape.z, + window_size.z, + stride.z, + q_tile_shape.z, + kv_tile_shape.z); +} + +} // namespace fna +} // namespace cuda +} // namespace natten diff --git a/natten/csrc/include/natten/cuda/fna_blackwell/collective/fna_common.hpp b/natten/csrc/include/natten/cuda/fna_blackwell/collective/fna_common.hpp new file mode 100644 index 00000000..128f3ced --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna_blackwell/collective/fna_common.hpp @@ -0,0 +1,180 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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::fna::collective { + +using namespace cute; + +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 < + class a_type, + class b_type, + class c_type, + int M, + int N, + UMMA::Major a_major, + UMMA::Major b_major, + UMMA::ScaleIn a_neg, + UMMA::ScaleIn b_neg, + class... TAs, + class... TMs> +CUTE_HOST_DEVICE constexpr auto to_tiled_mma_sm100_ts( + TiledMMA< + MMA_Atom< + MMA_Traits< + SM100_MMA_F8F6F4_SS, + a_type, + b_type, + c_type, + cute::C, + cute::C, + cute::integral_constant, + cute::integral_constant, + cute::integral_constant, + cute::integral_constant>, + TAs...>, + TMs...>) { + return TiledMMA< + MMA_Atom< + MMA_Traits>, + TAs...>, + TMs...>{}; +} + +template < + class a_type, + class b_type, + class c_type, + int M, + int N, + UMMA::Major a_major, + UMMA::Major b_major, + UMMA::ScaleIn a_neg, + UMMA::ScaleIn b_neg, + class... TAs, + class... TMs> +CUTE_HOST_DEVICE constexpr auto to_tiled_mma_sm100_ts(TiledMMA< + MMA_Atom< + SM100_MMA_F16BF16_SS< + a_type, + b_type, + c_type, + M, + N, + a_major, + b_major, + a_neg, + b_neg>, + TAs...>, + TMs...>) { + return TiledMMA< + MMA_Atom< + SM100_MMA_F16BF16_TS< + a_type, + b_type, + c_type, + M, + N, + a_major, + b_major, + a_neg, + b_neg, + UMMA::Saturate::False>, + 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::fna::collective diff --git a/natten/csrc/include/natten/cuda/fna_blackwell/collective/fna_fusion.hpp b/natten/csrc/include/natten/cuda/fna_blackwell/collective/fna_fusion.hpp new file mode 100644 index 00000000..ab1e31c5 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna_blackwell/collective/fna_fusion.hpp @@ -0,0 +1,609 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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::fna::collective { + +using namespace cute; + +// Tuple utils + +template +CUTE_HOST_DEVICE constexpr auto floor_tuple(T0 const& t0, T1 const& t1) { + return transform_leaf( + t0, t1, [&](auto const& a, auto const& b) { return (a / b) * b; }); +} + +template +CUTE_HOST_DEVICE constexpr auto ceil_tuple(T0 const& t0, T1 const& t1) { + return transform_leaf( + t0, t1, [&](auto const& a, auto const& b) { return ceil_div(a, b) * b; }); +} + +template +CUTE_HOST_DEVICE constexpr auto tuple_sub(T0 const& t0, T1 const& t1) { + return transform_leaf( + t0, t1, [&](auto const& a, auto const& b) { return a - b; }); +} + +template +CUTE_HOST_DEVICE constexpr auto tuple_add(T0 const& t0, T1 const& t1) { + return transform_leaf( + t0, t1, [&](auto const& a, auto const& b) { return a + b; }); +} + +template +CUTE_HOST_DEVICE constexpr auto tuple_mul(T0 const& t0, T1 const& t1) { + return transform_leaf( + t0, t1, [&](auto const& a, auto const& b) { return a * b; }); +} + +template +CUTE_HOST_DEVICE bool tuple_leq(T t0, T t1) { + static_assert(rank(T{}) > 0 && rank(T{}) < 4); + if constexpr (rank(T{}) == 1) { + return get<0>(t0) <= get<0>(t1); + } else if constexpr (rank(T{}) == 2) { + return get<0>(t0) <= get<0>(t1) && get<1>(t0) <= get<1>(t1); + } else { + return get<0>(t0) <= get<0>(t1) && get<1>(t0) <= get<1>(t1) && + get<2>(t0) <= get<2>(t1); + } +} + +// NA mask utils +template +CUTE_HOST_DEVICE int get_win_start( + int index, + int window_left, + int window_right, + int stride, + int length) { + if constexpr (IsCausal) { + int32_t stride_group_leader_idx = + cutlass::fast_min((index / stride) * stride + stride - 1, length - 1); + // window_size == window_left + window_right + 1 + // return cutlass::fast_max(stride_group_leader_idx - window_size + 1, 0); + return cutlass::fast_max( + stride_group_leader_idx - window_left - window_right, 0); + + } else { + auto stride_group_leader_idx = cutlass::fast_min( + ((index / stride) * stride) + (stride / 2), length - 1); + + return cutlass::fast_max(stride_group_leader_idx - window_left, 0) + + ((stride_group_leader_idx + window_right >= length) * + (length - window_right - stride_group_leader_idx - 1)); + } +} + +template +CUTE_HOST_DEVICE int get_win_end( + int index, + int start, + int window_size, + int length) { + if constexpr (IsCausal) { + return cutlass::fast_min(index + 1, length); + } else { + return start + window_size; + } +} + +template +CUTE_HOST_DEVICE auto get_window_start( + Coord index, + NADim window_left, + NADim window_right, + NADim stride, + NADim length) { + static_assert(rank(index) > 0 && rank(index) < 4); + static_assert(rank(index) == rank(Causal{})); + static_assert(rank(index) == rank(length)); + if constexpr (rank(index) == 1) { + return make_tuple(get_win_start(Causal{})>( + get<0>(index), + get<0>(window_left), + get<0>(window_right), + get<0>(stride), + get<0>(length))); + } else if constexpr (rank(index) == 2) { + return make_tuple( + get_win_start(Causal{})>( + get<0>(index), + get<0>(window_left), + get<0>(window_right), + get<0>(stride), + get<0>(length)), + get_win_start(Causal{})>( + get<1>(index), + get<1>(window_left), + get<1>(window_right), + get<1>(stride), + get<1>(length))); + } else { + return make_tuple( + get_win_start(Causal{})>( + get<0>(index), + get<0>(window_left), + get<0>(window_right), + get<0>(stride), + get<0>(length)), + get_win_start(Causal{})>( + get<1>(index), + get<1>(window_left), + get<1>(window_right), + get<1>(stride), + get<1>(length)), + get_win_start(Causal{})>( + get<2>(index), + get<2>(window_left), + get<2>(window_right), + get<2>(stride), + get<2>(length))); + } +} + +template +CUTE_HOST_DEVICE auto get_window_end( + Coord index, + NADim start, + NADim window_size, + NADim length) { + static_assert(rank(index) > 0 && rank(index) < 4); + static_assert(rank(index) == rank(Causal{})); + static_assert(rank(index) == rank(length)); + if constexpr (rank(index) == 1) { + return make_tuple(get_win_end(Causal{})>( + get<0>(index), get<0>(start), get<0>(window_size), get<0>(length))); + } else if constexpr (rank(index) == 2) { + return make_tuple( + get_win_end(Causal{})>( + get<0>(index), get<0>(start), get<0>(window_size), get<0>(length)), + get_win_end(Causal{})>( + get<1>(index), get<1>(start), get<1>(window_size), get<1>(length))); + } else { + return make_tuple( + get_win_end(Causal{})>( + get<0>(index), get<0>(start), get<0>(window_size), get<0>(length)), + get_win_end(Causal{})>( + get<1>(index), get<1>(start), get<1>(window_size), get<1>(length)), + get_win_end(Causal{})>( + get<2>(index), get<2>(start), get<2>(window_size), get<2>(length))); + } +} + +template +CUTE_HOST_DEVICE bool is_neighbor( + Coord kv_coord, + NADim kv_start, + NADim kv_end) { + static_assert(rank(kv_coord) > 0 && rank(kv_coord) < 4); + static_assert(rank(kv_coord) == rank(kv_start)); + if constexpr (rank(kv_coord) == 1) { + return get<0>(kv_coord) >= get<0>(kv_start) && + get<0>(kv_coord) < get<0>(kv_end); + } else if constexpr (rank(kv_coord) == 2) { + return get<0>(kv_coord) >= get<0>(kv_start) && + get<0>(kv_coord) < get<0>(kv_end) && + get<1>(kv_coord) >= get<1>(kv_start) && + get<1>(kv_coord) < get<1>(kv_end); + } else { + return get<0>(kv_coord) >= get<0>(kv_start) && + get<0>(kv_coord) < get<0>(kv_end) && + get<1>(kv_coord) >= get<1>(kv_start) && + get<1>(kv_coord) < get<1>(kv_end) && + get<2>(kv_coord) >= get<2>(kv_start) && + get<2>(kv_coord) < get<2>(kv_end); + } +} + +template +CUTE_HOST_DEVICE bool is_within_bounds(Coord kv_coord, NADim kv_shape) { + static_assert(rank(kv_coord) > 0 && rank(kv_coord) < 4); + static_assert(rank(kv_coord) == rank(kv_shape)); + if constexpr (rank(kv_coord) == 1) { + return get<0>(kv_coord) < get<0>(kv_shape); + } else if constexpr (rank(kv_coord) == 2) { + return get<0>(kv_coord) < get<0>(kv_shape) && + get<1>(kv_coord) < get<1>(kv_shape); + } else { + return get<0>(kv_coord) < get<0>(kv_shape) && + get<1>(kv_coord) < get<1>(kv_shape) && + get<2>(kv_coord) < get<2>(kv_shape); + } +} + +CUTE_HOST_DEVICE +int qkv_fix_dilation(int qkv_shape, int dilation, int dilation_group) { + auto padding = + 1 - ((dilation_group + (dilation - (qkv_shape % dilation))) / dilation); + return (qkv_shape / dilation) + padding; +} + +template +CUTE_HOST_DEVICE auto correct_qkv_shape_wrt_dilation( + NADim qkv_shape, + NADim dilation, + Coord dilation_group) { + static_assert(rank(qkv_shape) > 0 && rank(qkv_shape) < 4); + static_assert(rank(qkv_shape) == rank(dilation_group)); + if constexpr (rank(qkv_shape) == 1) { + return make_tuple(qkv_fix_dilation( + get<0>(qkv_shape), get<0>(dilation), get<0>(dilation_group))); + } else if constexpr (rank(qkv_shape) == 2) { + return make_tuple( + qkv_fix_dilation( + get<0>(qkv_shape), get<0>(dilation), get<0>(dilation_group)), + qkv_fix_dilation( + get<1>(qkv_shape), get<1>(dilation), get<1>(dilation_group))); + } else { + return make_tuple( + qkv_fix_dilation( + get<0>(qkv_shape), get<0>(dilation), get<0>(dilation_group)), + qkv_fix_dilation( + get<1>(qkv_shape), get<1>(dilation), get<1>(dilation_group)), + qkv_fix_dilation( + get<2>(qkv_shape), get<2>(dilation), get<2>(dilation_group))); + } +} + +template +struct NeighborhoodAttentionMask { + using Causal = Causal_; + static_assert(rank(Causal{}) >= 1 && rank(Causal{}) < 4); + + // QKV shape Correction + // Only if dilated and input size % dilation != 0 + // NOTE: every warp role has to execute this before using the mask!! + template + CUTLASS_DEVICE auto correct_qkv_shape( + ProblemShape const& problem_shape, + QKVShape const& qkv_shape, // this is pre-padding, pre-token permute, just + // the original shape of the sequence mode in + // the self attention + BlkCoord const& blk_coord, + Dilation const& dilation, + int num_dilation_groups) { + auto batch_idx = get<2, 1>(blk_coord); + + auto dilation_group_idx = batch_idx % num_dilation_groups; + auto dilation_group_crd = idx2crd(dilation_group_idx, dilation); + + return correct_qkv_shape_wrt_dilation( + qkv_shape, dilation, dilation_group_crd); + } + + template < + class BlkCoord, + class MultiDimTileShape, + class QKVShape, + class NAParams> + CUTLASS_DEVICE auto get_trip_count( + BlkCoord const& blk_coord, + MultiDimTileShape const& multi_dim_tile_shapes, + QKVShape const& q_shape, + QKVShape const& qkv_shape, + NAParams const& na_params) { + auto [q_tile_shape, kv_tile_shape] = multi_dim_tile_shapes; + + auto [window_size, window_left, window_right, stride] = na_params; + + auto q_tiled = ceil_div(q_shape, q_tile_shape); + + // Map query index back to the "original" coord + auto q_tile_coord = idx2crd(static_cast(get<0>(blk_coord)), q_tiled); + auto q_coord = tuple_mul(q_tile_coord, q_tile_shape); + + auto q_tile_offset_last = idx2crd(size(q_tile_shape) - 1, q_tile_shape); + auto q_coord_last = tuple_add(q_coord, q_tile_offset_last); + + auto kv_start_actual = get_window_start( + q_coord, window_left, window_right, stride, qkv_shape); + + auto last_kv_start_actual = get_window_start( + q_coord_last, window_left, window_right, stride, qkv_shape); + auto kv_end_actual = get_window_end( + q_coord_last, last_kv_start_actual, window_size, qkv_shape); + + auto kv_start = floor_tuple(kv_start_actual, kv_tile_shape); + auto kv_end = ceil_tuple(kv_end_actual, kv_tile_shape); + + auto kv_diff = tuple_sub(kv_end, kv_start); + auto kv_diff_tiles = ceil_div(kv_diff, kv_tile_shape); + + return make_tuple(kv_start, kv_diff_tiles); + } + + template + CUTLASS_DEVICE auto get_extra_trip_count( + MultiDimTileShape const& multi_dim_tile_shapes, + int num_extra_kv) { + auto [q_tile_shape, kv_tile_shape] = multi_dim_tile_shapes; + + int num_unmasked_trips = num_extra_kv / size(kv_tile_shape); + bool has_masked_trips = num_extra_kv % size(kv_tile_shape) != 0; + + return make_tuple(num_unmasked_trips, has_masked_trips); + } + + template + CUTLASS_DEVICE auto get_extra_trip_info( + MultiDimTileShape const& multi_dim_tile_shapes, + int num_extra_kv) { + auto [q_tile_shape, kv_tile_shape] = multi_dim_tile_shapes; + + auto [num_trips, has_masked_trips] = + get_extra_trip_count(multi_dim_tile_shapes, num_extra_kv); + if (has_masked_trips) { + num_trips++; + } + + return num_trips; + } + + template < + class AccQK, + class IndexQK, + class MultiDimTileShape, + class QKVShape, + class NAParams> + CUTLASS_DEVICE void apply_mask( + AccQK& acc_qk, + IndexQK const& index_qk, + MultiDimTileShape const& multi_dim_tile_shapes, + QKVShape const& q_shape, + QKVShape const& qkv_shape, + NAParams const& na_params, + QKVShape const& blk_kv_offset, + QKVShape const& kv_diff_tiles) { + auto [q_tile_shape, kv_tile_shape] = multi_dim_tile_shapes; + auto [window_size, window_left, window_right, stride] = na_params; + + auto q_tiled = ceil_div(q_shape, q_tile_shape); + + // NOTE: the following assumes each thread owns exactly 1 row, and + // contiguous columns. This is an assumption made in the FMHA fwd kernel. It + // also assumes the first kv_idx always evenly divides KV tile --size--. + // Otherwise we'd need a modulo like we have for q_idx. + auto [q_idx, kv_idx] = index_qk(0); + + int q_tile_idx = q_idx / size(q_tile_shape); + int q_tile_res = q_idx % size(q_tile_shape); + + auto q_tile_coord = idx2crd(q_tile_idx, q_tiled); + auto q_tile_offset = idx2crd(q_tile_res, q_tile_shape); + auto q_coord = + tuple_add(tuple_mul(q_tile_coord, q_tile_shape), q_tile_offset); + + auto kv_start = get_window_start( + q_coord, window_left, window_right, stride, qkv_shape); + auto kv_end = + get_window_end(q_coord, kv_start, window_size, qkv_shape); + + int kv_tile_idx = kv_idx / size(kv_tile_shape); + + auto kv_tile_coord = idx2crd(kv_tile_idx, kv_diff_tiles); + auto kv_tile_offset = + tuple_add(blk_kv_offset, tuple_mul(kv_tile_coord, kv_tile_shape)); + + auto kv_ctr = make_identity_tensor(kv_tile_shape); + auto kv_ctr_offset = domain_offset(kv_tile_offset, kv_ctr); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(acc_qk); i++) { + auto kv_coord = kv_ctr_offset(i); + + if (not is_neighbor(kv_coord, kv_start, kv_end)) { + acc_qk(i) = -INFINITY; + } + } + } + + template < + class AccQK, + class IndexQK, + class MultiDimTileShape, + class QKVShape, + class NAParams> + CUTLASS_DEVICE void apply_padded_mask( + AccQK& acc_qk, + IndexQK const& index_qk, + MultiDimTileShape const& multi_dim_tile_shapes, + QKVShape const& qkv_shape, + NAParams const& na_params, + QKVShape const& blk_kv_offset, + QKVShape const& kv_diff_tiles) { + auto [q_tile_shape, kv_tile_shape] = multi_dim_tile_shapes; + + // NOTE: the following assumes each contiguous columns. + // This is an assumption made in the FMHA fwd kernel. + // It also assumes the first kv_idx always evenly divides KV tile --size--. + // Otherwise we'd need a modulo like we have for q_idx. + auto [q_idx, kv_idx] = index_qk(0); + + int kv_tile_idx = kv_idx / size(kv_tile_shape); + + auto kv_tile_coord = idx2crd(kv_tile_idx, kv_diff_tiles); + auto kv_tile_offset = + tuple_add(blk_kv_offset, tuple_mul(kv_tile_coord, kv_tile_shape)); + + auto kv_ctr = make_identity_tensor(kv_tile_shape); + auto kv_ctr_offset = domain_offset(kv_tile_offset, kv_ctr); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(acc_qk); i++) { + auto kv_coord = kv_ctr_offset(i); + + if (not is_within_bounds(kv_coord, qkv_shape)) { + acc_qk(i) = -INFINITY; + } + } + } + + template + CUTLASS_DEVICE void apply_extra_kv_mask( + AccQK& acc_qk, + IndexQK const& index_qk, + int num_extra_kv) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(acc_qk); i++) { + if (get<1>(index_qk(i)) >= num_extra_kv) { + acc_qk(i) = -INFINITY; + } + } + } +}; + +// Misc +template +CUTE_HOST_DEVICE constexpr auto get_window_left(NADim const& window_size) { + return transform_leaf(window_size, [&](auto const& w) { return w / 2; }); +} + +template +CUTE_HOST_DEVICE constexpr auto get_window_right(NADim const& window_size) { + return transform_leaf( + window_size, [&](auto const& w) { return (w / 2) + ((w % 2) - 1); }); +} + +CUTE_HOST_DEVICE +bool is_fully_block_sparse_int( + int input_size, + int window_size, + int stride, + int q_tile_size, + int kv_tile_size) { + int num_stride_groups = ceil_div(input_size, stride); + int window_left = window_size / 2; + int window_right = (window_size / 2) + ((window_size % 2) - 1); + int stride_group_center = stride / 2; + + int last_stride_group_center = + min((input_size / stride) * stride + stride_group_center, input_size - 1); + int last_stride_group_window_start = get_win_start( + last_stride_group_center, window_left, window_right, stride, input_size); + + return input_size == window_size || + (stride % q_tile_size == 0 && window_size % kv_tile_size == 0 && + (input_size % stride == 0 || + last_stride_group_window_start % kv_tile_size == 0) && + (((window_left - stride_group_center) % kv_tile_size == 0) || + ( + // input_size % stride == 0 && + input_size % kv_tile_size == 0 && num_stride_groups == 2 && + stride + stride_group_center >= input_size - window_size))); +} + +template +CUTE_HOST_DEVICE bool fully_block_sparse( + NADim input_size, + NADim window_size, + NADim stride, + QTile q_tile_shape, + KVTile kv_tile_shape) { +#ifdef NATTEN_DISABLE_FULLY_BLOCK_SPARSE_FAST_PATH + return false; +#else + // Causal masking can never be fully block-sparse (unless tile/block size is + // 1) + // TODO? + if (cute::any_of(Causal{}, [&](auto const& a) { return a; })) { + return false; + } + + static_assert(rank(input_size) > 0 && rank(input_size) < 4); + static_assert(rank(q_tile_shape) == rank(input_size)); + static_assert(rank(kv_tile_shape) == rank(input_size)); + + if constexpr (rank(input_size) == 1) { + return is_fully_block_sparse_int( + get<0>(input_size), + get<0>(window_size), + get<0>(stride), + get<0>(q_tile_shape), + get<0>(kv_tile_shape)); + } else if constexpr (rank(input_size) == 2) { + return is_fully_block_sparse_int( + get<0>(input_size), + get<0>(window_size), + get<0>(stride), + get<0>(q_tile_shape), + get<0>(kv_tile_shape)) && + is_fully_block_sparse_int( + get<1>(input_size), + get<1>(window_size), + get<1>(stride), + get<1>(q_tile_shape), + get<1>(kv_tile_shape)); + } else { + return is_fully_block_sparse_int( + get<0>(input_size), + get<0>(window_size), + get<0>(stride), + get<0>(q_tile_shape), + get<0>(kv_tile_shape)) && + is_fully_block_sparse_int( + get<1>(input_size), + get<1>(window_size), + get<1>(stride), + get<1>(q_tile_shape), + get<1>(kv_tile_shape)) && + is_fully_block_sparse_int( + get<2>(input_size), + get<2>(window_size), + get<2>(stride), + get<2>(q_tile_shape), + get<2>(kv_tile_shape)); + } +#endif +} + +template +CUTE_HOST_DEVICE bool is_dilated(NADim dilation) { + static_assert(rank(dilation) > 0 && rank(dilation) < 4); + + if constexpr (rank(dilation) == 1) { + return get<0>(dilation) != 1; + } else if constexpr (rank(dilation) == 2) { + return get<0>(dilation) != 1 || get<1>(dilation) != 1; + } else { + return get<0>(dilation) != 1 || get<1>(dilation) != 1 || + get<2>(dilation) != 1; + } +} + +} // namespace cutlass::fna::collective diff --git a/natten/csrc/include/natten/cuda/fna_blackwell/collective/fna_fusion_bwd.hpp b/natten/csrc/include/natten/cuda/fna_blackwell/collective/fna_fusion_bwd.hpp new file mode 100644 index 00000000..140ff1a8 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna_blackwell/collective/fna_fusion_bwd.hpp @@ -0,0 +1,450 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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" + +#include "natten/cuda/fna_blackwell/collective/fna_fusion.hpp" + +namespace cutlass::fna::collective { + +using namespace cute; + +// BWD + +// BWD mask +template +CUTE_HOST_DEVICE int get_bwd_win_start( + int index, + int stride_group_offset, + int window_left, + int window_right, + int window_size, + int stride, + int length) { + if constexpr (IsCausal) { + return index; + } else { + auto window_start_pre_trunc = + (index >= window_size) * (index - window_right + stride_group_offset); + return (window_start_pre_trunc / stride) * stride; + } +} + +template +CUTE_HOST_DEVICE int get_bwd_win_end( + int index, + int stride_group_offset, + int window_left, + int window_right, + int window_size, + int stride, + int length) { + if constexpr (IsCausal) { + // Window end is always the last query that attends to this key/value + 1 + // In the strided case, it will always be the last query in a stride + // group + 1, which means it will always be a multiple of stride. + // In the case of incomplete stridegroups (input size % stride != 0), + // we must correct backward window end by incrementing first non-attending + // stridegroup by 1. + auto first_non_attending_stride_group_idx = + (index + window_size) / stride + (index + window_size >= length); + return cutlass::fast_min( + first_non_attending_stride_group_idx * stride, length); + + } else { + auto bound_check = (index >= (length - window_size)); + auto window_end = + ((index + window_left + stride_group_offset + 1) / stride) * stride; + return (bound_check * length) + ((1 - bound_check) * window_end); + } +} + +template +CUTE_HOST_DEVICE auto get_bwd_window_start( + NADim index, + NADim stride_group_offset, + NADim window_left, + NADim window_right, + NADim window_size, + NADim stride, + NADim length) { + static_assert(rank(index) > 0 && rank(index) < 4); + static_assert(rank(index) == rank(Causal{})); + if constexpr (rank(index) == 1) { + return make_tuple(get_bwd_win_start(Causal{})>( + get<0>(index), + get<0>(stride_group_offset), + get<0>(window_left), + get<0>(window_right), + get<0>(window_size), + get<0>(stride), + get<0>(length))); + } else if constexpr (rank(index) == 2) { + return make_tuple( + get_bwd_win_start(Causal{})>( + get<0>(index), + get<0>(stride_group_offset), + get<0>(window_left), + get<0>(window_right), + get<0>(window_size), + get<0>(stride), + get<0>(length)), + get_bwd_win_start(Causal{})>( + get<1>(index), + get<1>(stride_group_offset), + get<1>(window_left), + get<1>(window_right), + get<1>(window_size), + get<1>(stride), + get<1>(length))); + } else { + return make_tuple( + get_bwd_win_start(Causal{})>( + get<0>(index), + get<0>(stride_group_offset), + get<0>(window_left), + get<0>(window_right), + get<0>(window_size), + get<0>(stride), + get<0>(length)), + get_bwd_win_start(Causal{})>( + get<1>(index), + get<1>(stride_group_offset), + get<1>(window_left), + get<1>(window_right), + get<1>(window_size), + get<1>(stride), + get<1>(length)), + get_bwd_win_start(Causal{})>( + get<2>(index), + get<2>(stride_group_offset), + get<2>(window_left), + get<2>(window_right), + get<2>(window_size), + get<2>(stride), + get<2>(length))); + } +} + +template +CUTE_HOST_DEVICE auto get_bwd_window_end( + NADim index, + NADim stride_group_offset, + NADim window_left, + NADim window_right, + NADim window_size, + NADim stride, + NADim length) { + static_assert(rank(index) > 0 && rank(index) < 4); + static_assert(rank(index) == rank(Causal{})); + if constexpr (rank(index) == 1) { + return make_tuple(get_bwd_win_end(Causal{})>( + get<0>(index), + get<0>(stride_group_offset), + get<0>(window_left), + get<0>(window_right), + get<0>(window_size), + get<0>(stride), + get<0>(length))); + } else if constexpr (rank(index) == 2) { + return make_tuple( + get_bwd_win_end(Causal{})>( + get<0>(index), + get<0>(stride_group_offset), + get<0>(window_left), + get<0>(window_right), + get<0>(window_size), + get<0>(stride), + get<0>(length)), + get_bwd_win_end(Causal{})>( + get<1>(index), + get<1>(stride_group_offset), + get<1>(window_left), + get<1>(window_right), + get<1>(window_size), + get<1>(stride), + get<1>(length))); + } else { + return make_tuple( + get_bwd_win_end(Causal{})>( + get<0>(index), + get<0>(stride_group_offset), + get<0>(window_left), + get<0>(window_right), + get<0>(window_size), + get<0>(stride), + get<0>(length)), + get_bwd_win_end(Causal{})>( + get<1>(index), + get<1>(stride_group_offset), + get<1>(window_left), + get<1>(window_right), + get<1>(window_size), + get<1>(stride), + get<1>(length)), + get_bwd_win_end(Causal{})>( + get<2>(index), + get<2>(stride_group_offset), + get<2>(window_left), + get<2>(window_right), + get<2>(window_size), + get<2>(stride), + get<2>(length))); + } +} + +// Backward pass mask +// Does not allow extra KV fusion +template +struct NeighborhoodAttentionBackwardMask { + using Causal = Causal_; + static_assert(rank(Causal{}) >= 1 && rank(Causal{}) < 4); + + // Identical to NeighborhoodAttentionMask::correct_qkv_shape + // QKV shape Correction + // Only if dilated and input size % dilation != 0 + // NOTE: every warp role has to execute this before using the mask!! + template + CUTLASS_DEVICE auto correct_qkv_shape( + ProblemShape const& problem_shape, + QKVShape const& qkv_shape, // this is pre-padding, pre-token permute, just + // the original shape of the sequence mode in + // the self attention + BlkCoord const& blk_coord, + Dilation const& dilation, + int num_dilation_groups) { + auto batch_idx = get<4, 1>(blk_coord); + auto dilation_group_idx = batch_idx % num_dilation_groups; + + auto dilation_group_crd = idx2crd(dilation_group_idx, dilation); + return correct_qkv_shape_wrt_dilation( + qkv_shape, dilation, dilation_group_crd); + } + + // Unlike in forward pass, trip counts are over Q tiles and not KV tiles. + template < + class BlkCoord, + class MultiDimTileShape, + class QKVShape, + class NAParams> + CUTLASS_DEVICE auto get_trip_count( + BlkCoord const& blk_coord, + MultiDimTileShape const& multi_dim_tile_shapes, + QKVShape const& kv_shape, + QKVShape const& qkv_shape, + NAParams const& na_params) { + auto [q_tile_shape, kv_tile_shape] = multi_dim_tile_shapes; + + auto [window_size, window_left, window_right, stride, stride_group_offset] = + na_params; + + auto kv_tiled = ceil_div(kv_shape, kv_tile_shape); + + // Map KV index back to coord + auto kv_tile_coord = idx2crd(static_cast(get<1>(blk_coord)), kv_tiled); + auto kv_coord = tuple_mul(kv_tile_coord, kv_tile_shape); + + auto kv_tile_offset_last = idx2crd(size(kv_tile_shape) - 1, kv_tile_shape); + auto kv_coord_last = tuple_add(kv_coord, kv_tile_offset_last); + + // q start and end instead of kv like in forward pass + auto q_start_actual = get_bwd_window_start( + kv_coord, + stride_group_offset, + window_left, + window_right, + window_size, + stride, + qkv_shape); + + auto last_q_start_actual = get_bwd_window_start( + kv_coord_last, + stride_group_offset, + window_left, + window_right, + window_size, + stride, + qkv_shape); + auto q_end_actual = get_bwd_window_end( + kv_coord_last, + stride_group_offset, + window_left, + window_right, + window_size, + stride, + qkv_shape); + + auto q_start = floor_tuple(q_start_actual, q_tile_shape); + auto q_end = ceil_tuple(q_end_actual, q_tile_shape); + + auto q_diff = tuple_sub(q_end, q_start); + auto q_diff_tiles = ceil_div(q_diff, q_tile_shape); + + return make_tuple(q_start, q_diff_tiles); + } + + template < + class AccQK, + class IndexQ, + class IndexQK, + class MultiDimTileShape, + class QKVShape, + class NAParams> + CUTLASS_DEVICE void apply_mask( + AccQK& acc_qk, + IndexQ const& index_q, + IndexQK const& index_qk, + MultiDimTileShape const& multi_dim_tile_shapes, + QKVShape const& kv_shape, + QKVShape const& qkv_shape, + NAParams const& na_params, + QKVShape const& blk_q_offset, + QKVShape const& q_diff_tiles) { + auto [q_tile_shape, kv_tile_shape] = multi_dim_tile_shapes; + auto [window_size, window_left, window_right, stride, stride_group_offset] = + na_params; + + auto kv_tiled = ceil_div(kv_shape, kv_tile_shape); + + // NOTE: Unlike the forward pass, there is no guarantee that each thread + // owns entire rows, and starts with the first dot product in the tile + // (which eliminated the need to do an extra modulo). However, each thread + // visits only one row (one key, many queries), which simplifies things + // again like in the forward pass, compared to Hopper. We can just compute + // window start and end once, and just do bound checks for each dot product. + auto [q_idx, kv_idx] = index_qk(0); + + int kv_tile_idx = kv_idx / size(kv_tile_shape); + int kv_tile_res = kv_idx % size(kv_tile_shape); + + auto kv_tile_coord = idx2crd(kv_tile_idx, kv_tiled); + auto kv_tile_offset = idx2crd(kv_tile_res, kv_tile_shape); + auto kv_coord = + tuple_add(tuple_mul(kv_tile_coord, kv_tile_shape), kv_tile_offset); + + auto q_start = get_bwd_window_start( + kv_coord, + stride_group_offset, + window_left, + window_right, + window_size, + stride, + qkv_shape); + auto q_end = get_bwd_window_end( + kv_coord, + stride_group_offset, + window_left, + window_right, + window_size, + stride, + qkv_shape); + + int q_tile_idx = q_idx / size(q_tile_shape); + int q_tile_res = q_idx % size(q_tile_shape); + + auto q_tile_coord = idx2crd(q_tile_idx, q_diff_tiles); + auto q_iter_offset = idx2crd(q_tile_res, q_tile_shape); + auto q_tile_offset = tuple_add( + q_iter_offset, + tuple_add(blk_q_offset, tuple_mul(q_tile_coord, q_tile_shape))); + + auto q_ctr = make_identity_tensor(q_tile_shape); + auto q_ctr_offset = domain_offset(q_tile_offset, q_ctr); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(acc_qk); i++) { + auto q_coord = q_ctr_offset(index_q(i) - q_idx); + + if (not is_neighbor(q_coord, q_start, q_end)) { + acc_qk(i) = -INFINITY; + } + } + } + + template < + class AccQK, + class IndexQ, + class MultiDimTileShape, + class QKVShape, + class NAParams> + CUTLASS_DEVICE void apply_padded_mask( + AccQK& acc_qk, + IndexQ const& index_q, + MultiDimTileShape const& multi_dim_tile_shapes, + QKVShape const& qkv_shape, + NAParams const& na_params, + QKVShape const& blk_q_offset, + QKVShape const& q_diff_tiles) { + auto [q_tile_shape, kv_tile_shape] = multi_dim_tile_shapes; + + // NOTE: Unlike the forward pass, there is no guarantee that each thread + // owns entire rows, and starts with the first dot product in the tile + // (which eliminated the need to do an extra modulo). However, each thread + // visits only one row (one key, many queries), which simplifies things + // again like in the forward pass, compared to Hopper. We can just compute + // window start and end once, and just do bound checks for each dot product. + auto q_idx = index_q(0); + + int q_tile_idx = q_idx / size(q_tile_shape); + int q_tile_res = q_idx % size(q_tile_shape); + + auto q_tile_coord = idx2crd(q_tile_idx, q_diff_tiles); + auto q_iter_offset = idx2crd(q_tile_res, q_tile_shape); + auto q_tile_offset = tuple_add( + q_iter_offset, + tuple_add(blk_q_offset, tuple_mul(q_tile_coord, q_tile_shape))); + + auto q_ctr = make_identity_tensor(q_tile_shape); + auto q_ctr_offset = domain_offset(q_tile_offset, q_ctr); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(acc_qk); i++) { + auto q_coord = q_ctr_offset(index_q(i) - q_idx); + + if (not is_within_bounds(q_coord, qkv_shape)) { + acc_qk(i) = -INFINITY; + } + } + } +}; + +// Misc + +template +CUTE_HOST_DEVICE constexpr auto get_bwd_stride_offset(NADim const& stride) { + return transform_leaf( + stride, [&](auto const& s) { return (s - (s / 2) - 1); }); +} + +} // namespace cutlass::fna::collective diff --git a/natten/csrc/include/natten/cuda/fna_blackwell/collective/sm100_fna_fwd_epilogue_tma_warpspecialized.hpp b/natten/csrc/include/natten/cuda/fna_blackwell/collective/sm100_fna_fwd_epilogue_tma_warpspecialized.hpp new file mode 100644 index 00000000..590bd23f --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna_blackwell/collective/sm100_fna_fwd_epilogue_tma_warpspecialized.hpp @@ -0,0 +1,215 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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" + +namespace cutlass::fna::collective { + +template < + class Element, + class ElementAcc, + class TileShape, // Q, D, _ + class StrideO, // Q, D, B + class StrideLSE_, // Q, B + class OrderLoadEpilogue = cute::false_type> +struct Sm100FnaFwdEpilogueTmaWarpspecialized { + using Pipeline = cutlass::PipelineAsync<2>; + + // 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; + using StrideLSE = StrideLSE_; + using ElementOut = Element; + + static const int NumWarpsEpilogue = 1; + static const int NumWarpsLoad = 1; + + struct TensorStorage { + using SmemLayoutO = SmemLayoutO_; + cute::array_aligned> smem_o; + }; + + struct Arguments { + Element* ptr_O; + StrideO dO; + + ElementAcc* ptr_LSE; + StrideLSE dLSE; + }; + + 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; + + ElementAcc* ptr_LSE; + StrideLSE dLSE; + }; + + // FMHA and MLA have different input ProblemShapes; + // get problem_shape_O according to the input ProblemShape. + template + CUTLASS_DEVICE static constexpr auto get_problem_shape_O( + ProblemShape const& problem_shape) { + if constexpr (rank_v(ProblemShape{}))> == 2) { + return replace<1>( + select<0, 2, 3>(problem_shape), get<2, 0>(problem_shape)); + } else { + return select<0, 2, 3>(problem_shape); + } + } + + template + static Params to_underlying_arguments( + ProblemShape const& problem_shape, + Arguments const& args, + void* workspace = nullptr) { + auto ptr_O = args.ptr_O; + StrideO dO = args.dO; + + auto problem_shape_O = get_problem_shape_O(problem_shape); + + auto tma_store_o = make_tma_copy( + SM90_TMA_STORE{}, + make_tensor(ptr_O, problem_shape_O, dO), + SmemLayoutO{}(_, _, _0{})); + + return {tma_store_o, args.ptr_LSE, args.dLSE}; + } + + 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 Sm100FnaFwdEpilogueTmaWarpspecialized(const Params& params) + : params(params) {} + + 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) { + BlkCoord blk_coord = blk_coord_in; + 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; + + Tensor mO_qdl_p = + params.tma_store_o.get_tma_tensor(get_problem_shape_O(problem_shape)); + // offset mode 0 by (max_length - real_length) + // offset mode 3,1 by cumulative_length + real_length + // the ptr is already offset by - max_length + // so in total this achieves + int offs_0 = 0; + int offs_2_1 = 0; + + Tensor mO_qdl = domain_offset( + make_coord(offs_0, _0{}, make_coord(_0{}, offs_2_1)), mO_qdl_p); + + Tensor gO_qdl = + local_tile(mO_qdl, TileShape{}, make_coord(_, _, _), Step<_1, _1, X>{}); + Tensor gO = gO_qdl(_, _, _, _0{}, get<2>(blk_coord)); + 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>(); + + if constexpr (cute::is_same_v) { + cutlass::arch::NamedBarrier::arrive( + (NumWarpsLoad + NumWarpsEpilogue) * NumThreadsPerWarp, + cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); + } + + pipeline.consumer_release(pipeline_release_state); + ++pipeline_release_state; + } +}; + +} // namespace cutlass::fna::collective diff --git a/natten/csrc/include/natten/cuda/fna_blackwell/collective/sm100_fna_fwd_mainloop_tma_warpspecialized.hpp b/natten/csrc/include/natten/cuda/fna_blackwell/collective/sm100_fna_fwd_mainloop_tma_warpspecialized.hpp new file mode 100644 index 00000000..0053c6fc --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna_blackwell/collective/sm100_fna_fwd_mainloop_tma_warpspecialized.hpp @@ -0,0 +1,1454 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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 "natten/cuda/fna_blackwell/collective/fna_common.hpp" +#include "natten/cuda/fna_blackwell/collective/fna_fusion.hpp" +#include "natten/cuda/fna_blackwell/collective/sm100_fna_load_tma_warpspecialized.hpp" + +namespace cutlass::fna::collective { + +using namespace cute; + +template < + class Element_, + class ElementQK_, + class ElementPV_, + class TileShape_, + class StrideQ_, + class StrideK_, + class StrideV_, + class Mask_, + class QTileShape_, + class KVTileShape_, + class NADim_, + // shape here is QG K H + // and referes to the two softmax warps + // (2, 1, 1) means that they are stacked (best for large Q since it loads + // the least K/V) (1, 2, 1) means they sit side by side (best for small Q / + // large K) + class ThreadShape = Shape<_2, _1, _1>, + // Since shared memory is sufficient for FMHA, there is no need to reuse + // shared memory. + class OrderLoadEpilogue = cute::false_type> +struct Sm100FnaFwdMainloopTmaWarpspecialized { + using Element = Element_; + using ElementQK = ElementQK_; + using ElementPV = ElementPV_; + using TileShape = TileShape_; + using StrideQ = StrideQ_; + using StrideK = StrideK_; + using StrideV = StrideV_; + using Mask = Mask_; + + using QTileShape = QTileShape_; + using KVTileShape = KVTileShape_; + using MultiDimTileShape = cute::tuple; + + static_assert( + size(QTileShape{}) == get<0>(TileShape{}), + "QTileShape doesn't match the size of Q tile in the FMHA kernel."); + static_assert( + size(KVTileShape{}) == get<1>(TileShape{}), + "QTileShape doesn't match the size of Q tile in the FMHA kernel."); + + // (ahassani): realized too late that I made this assumption + static_assert(get<1>(ThreadShape{}) == 1 && get<2>(ThreadShape{}) == 1); + + static constexpr int StageCountQ = 2; + static constexpr int StageCountKV = sizeof(Element_) == 1 ? 4 : 3; + + 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{})); + + // Reuse shared memory for V and O. + static constexpr bool IsOrderLoadEpilogue = + std::is_same_v; + 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::PipelineTmaUmmaAsync< + StageCountQ, + typename CollectiveMmaQK::AtomThrShapeMNK>; + + // from load to mma warp, protects k/v in smem + using PipelineKV = cutlass::PipelineTmaUmmaAsync< + StageCountKV, + typename CollectiveMmaQK::AtomThrShapeMNK>; + + // 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( + TransactionBytesLoadK == TransactionBytesLoadV, + "K and V smem layouts must be of equal size"); + + using Load = Sm100FnaLoadTmaWarpspecialized< + Element, + StrideQ, + StrideK, + StrideV, + CollectiveMmaQK, + CollectiveMmaPV, + SmemLayoutQ, + SmemLayoutK, + SmemLayoutV, + TensorStorage, + PipelineQ, + PipelineKV, + Mask, + TileShape, + QTileShape, + KVTileShape, + NADim_>; + + 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 evenly_divides(args.load.q_shape, QTileShape{}) && + evenly_divides(args.load.kv_shape, KVTileShape{}) && + evenly_divides( + size<3, 1>(problem_shape), + size(args.load.dilation)) && // dilation groups are + // folded into batch + tuple_leq(args.load.window_size, args.load.qkv_shape) && + // TODO: check window size * dilation <= qkv shape + tuple_leq(args.load.stride, args.load.window_size); + } + + 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; + + auto qkv_shape = params.load.qkv_shape; + if (params.load.requires_qkv_fixup) { + qkv_shape = Mask{}.correct_qkv_shape( + problem_shape, + params.load.qkv_shape, + blk_coord, + params.load.dilation, + params.load.num_dilation_groups); + } else if (params.load.is_dilated) { + qkv_shape = ceil_div(params.load.qkv_shape, params.load.dilation); + } + + int mask_tile_count = size(get<1>(Mask{}.get_trip_count( + blk_coord, + MultiDimTileShape{}, + params.load.q_shape, + qkv_shape, + params.load.na_params))); + + 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 < + bool need_apply_mask, + class Stage, + class BlkCoord, + class CoordTensor, + class ProblemShape, + class NADim> + CUTLASS_DEVICE auto softmax_step( + NADim const& qkv_shape, + bool is_fully_block_sparse, + bool has_kv_padding, + NADim const& kv_offset, + NADim const& kv_tiles, + float& row_max, + float& row_sum, + Stage stage, + bool final_call, + BlkCoord const& blk_coord, + CoordTensor 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 = + size<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) { + if (not is_fully_block_sparse) { + Mask{}.apply_mask( + tTMEM_LOADrS, + tTMEM_LOADcS, + MultiDimTileShape{}, + params.load.q_shape, + qkv_shape, + params.load.na_params, + kv_offset, + kv_tiles); + } else if (has_kv_padding) { + Mask{}.apply_padded_mask( + tTMEM_LOADrS, + tTMEM_LOADcS, + MultiDimTileShape{}, + qkv_shape, + params.load.na_params, + kv_offset, + kv_tiles); + } + } + + 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 = (old_row_max == row_max_safe) + ? 0.5f + : 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) { + auto qkv_shape = params.load.qkv_shape; + bool is_fully_block_sparse = params.load.is_fully_block_sparse; + bool has_kv_padding = params.load.has_kv_padding; + if (params.load.requires_qkv_fixup) { + qkv_shape = Mask{}.correct_qkv_shape( + problem_shape, + params.load.qkv_shape, + blk_coord, + params.load.dilation, + params.load.num_dilation_groups); + is_fully_block_sparse = fully_block_sparse( + qkv_shape, + get<0>(params.load.na_params), + get<3>(params.load.na_params), + QTileShape{}, + KVTileShape{}); + has_kv_padding = not evenly_divides(qkv_shape, KVTileShape{}); + } else if (params.load.is_dilated) { + qkv_shape = ceil_div(params.load.qkv_shape, params.load.dilation); + is_fully_block_sparse = fully_block_sparse( + qkv_shape, + get<0>(params.load.na_params), + get<3>(params.load.na_params), + QTileShape{}, + KVTileShape{}); + has_kv_padding = not evenly_divides(qkv_shape, KVTileShape{}); + } + + auto [kv_offset, num_tiles] = Mask{}.get_trip_count( + blk_coord, + MultiDimTileShape{}, + params.load.q_shape, + qkv_shape, + params.load.na_params); + int mask_tile_count = size(num_tiles); + + 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); + + // (Self/neighborhood) attention iterations + CUTLASS_PRAGMA_NO_UNROLL + for (; mask_tile_count > 0; mask_tile_count -= 1) { + softmax_step( + qkv_shape, + is_fully_block_sparse, + has_kv_padding, + kv_offset, + num_tiles, + 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< + kCorrectionTileSize == 32, + SM100_TMEM_LOAD_32dp32b32x, + SM100_TMEM_LOAD_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 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<2>(TileShape{}) / 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<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))); + + if (scale != 1.0f) { + 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 < + class BlkCoord, + class ProblemShape, + class TensorStorageEpi, + class CollectiveEpilogue> + 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, + CollectiveEpilogue& epilogue) { + auto qkv_shape = params.load.qkv_shape; + if (params.load.requires_qkv_fixup) { + qkv_shape = Mask{}.correct_qkv_shape( + problem_shape, + params.load.qkv_shape, + blk_coord, + params.load.dilation, + params.load.num_dilation_groups); + } else if (params.load.is_dilated) { + qkv_shape = ceil_div(params.load.qkv_shape, params.load.dilation); + } + + int mask_tile_count = size(get<1>(Mask{}.get_trip_count( + blk_coord, + MultiDimTileShape{}, + params.load.q_shape, + qkv_shape, + params.load.na_params))); + + 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 = + (tTMEM_LOADVrS(kIdxOldRowMax) == tTMEM_LOADVrS(kIdxNewRowMax)) + ? 1.0f + : ::exp2f( + params.scale_softmax_log2 * + (tTMEM_LOADVrS(kIdxOldRowMax) - tTMEM_LOADVrS(kIdxNewRowMax))); + + pipeline_o.consumer_wait(pipeline_o_consumer_state); + + bool warp_do_correction = __any_sync(0xFFFFFFFF, scale != 1.0f); + if (warp_do_correction) { + 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 = (tTMEM_LOADVrS(kIdxOldRowMax) == tTMEM_LOADVrS(kIdxNewRowMax)) + ? 1.0f + : ::exp2f( + params.scale_softmax_log2 * + (tTMEM_LOADVrS(kIdxOldRowMax) - tTMEM_LOADVrS(kIdxNewRowMax))); + + pipeline_o.consumer_wait(pipeline_o_consumer_state); + + warp_do_correction = __any_sync(0xFFFFFFFF, scale != 1.0f); + if (warp_do_correction) { + 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), + select<0, 3>(problem_shape), + epilogue.params.dLSE); + + correction_epilogue( + params.scale_output / tTMEM_LOADVrS(kIdxFinalRowSum), _0{}, sO); + + if (epilogue.params.ptr_LSE != nullptr) { + int row_idx = + get<0>(tTMEM_LOADVcS(_0{})) + get<0>(TileShape{}) * get<0>(blk_coord); + + ElementPV lse = cutlass::fast_log(tTMEM_LOADVrS(kIdxFinalRowSum)) + + params.scale_softmax * tTMEM_LOADVrS(kIdxFinalRowMax); + + if (row_idx < get<0>(problem_shape)) { + gLSE(row_idx, get<2>(blk_coord)) = lse; + } + } + + 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 row_idx = get<0>(tTMEM_LOADVcS(_0{})) + + get<0>(TileShape{}) * get<0>(blk_coord) + get<0>(TileShapeQK{}); + + ElementPV lse = cutlass::fast_log(tTMEM_LOADVrS(kIdxFinalRowSum)) + + params.scale_softmax * tTMEM_LOADVrS(kIdxFinalRowMax); + + if (row_idx < get<0>(problem_shape)) { + gLSE(row_idx, get<2>(blk_coord)) = lse; + } + } + + 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 < + class BlkCoord, + class ProblemShape, + class ParamsProblemShape, + class TensorStorageEpi, + class CollectiveEpilogue> + 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), + select<0, 3>(problem_shape), + epilogue.params.dLSE); + float lse = -INFINITY; + int thread_idx = threadIdx.x % (4 * NumThreadsPerWarp); + +#define DSHOW(x) \ + print(#x ": "); \ + print(x); \ + print("\n") + if (threadIdx.x % 128 == 0 && block0()) { + DSHOW(sO); + } +#if 1 + + 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{})); +#endif + + if (epilogue.params.ptr_LSE != nullptr) { + int row_idx = thread_idx + get<0>(TileShape{}) * get<0>(blk_coord); + + if (row_idx < get<0>(problem_shape)) { + gLSE(row_idx, get<2>(blk_coord)) = lse; + } + } + + 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 row_idx = thread_idx + get<0>(TileShape{}) * get<0>(blk_coord) + + get<0>(TileShapeQK{}); + + if (row_idx < get<0>(problem_shape)) { + gLSE(row_idx, get<2>(blk_coord)) = lse; + } + } + + cutlass::arch::fence_view_async_shared(); + pipeline_epi.producer_commit(pipeline_epi_producer_state); + ++pipeline_epi_producer_state; + } +}; + +} // namespace cutlass::fna::collective diff --git a/natten/csrc/include/natten/cuda/fna_blackwell/collective/sm100_fna_load_tma_warpspecialized.hpp b/natten/csrc/include/natten/cuda/fna_blackwell/collective/sm100_fna_load_tma_warpspecialized.hpp new file mode 100644 index 00000000..6e6a0fae --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna_blackwell/collective/sm100_fna_load_tma_warpspecialized.hpp @@ -0,0 +1,393 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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 "cutlass/arch/memory_sm80.h" +#include "cutlass/cutlass.h" +#include "cutlass/gemm/collective/collective_builder.hpp" + +#include "natten/cuda/fna_blackwell/collective/fna_common.hpp" +#include "natten/cuda/fna_blackwell/collective/fna_fusion.hpp" + +namespace cutlass::fna::collective { + +using namespace cute; + +template < + class Element, + class StrideQ, + class StrideK, + class StrideV, + class CollectiveMmaQK, + class CollectiveMmaPV, + class SmemLayoutQ, + class SmemLayoutK, + class SmemLayoutV, + class TensorStorage, + class PipelineQ, + class PipelineKV, + class Mask, + class TileShape, + class QTileShape, + class KVTileShape, + class NADim> +struct Sm100FnaLoadTmaWarpspecialized { + using MultiDimTileShape = cute::tuple; + + static_assert( + size(QTileShape{}) == get<0>(TileShape{}), + "QTileShape doesn't match the size of Q tile in the FMHA kernel."); + static_assert( + size(KVTileShape{}) == get<1>(TileShape{}), + "QTileShape doesn't match the size of Q tile in the FMHA kernel."); + + using TileShapeQK = typename CollectiveMmaQK::TileShape; + using TileShapePV = typename CollectiveMmaPV::TileShape; + + struct Arguments { + const Element* ptr_Q; + StrideQ dQ; + const Element* ptr_K; + StrideK dK; + const Element* ptr_V; + StrideV dV; + NADim q_shape; + NADim kv_shape; + NADim qkv_shape; + NADim window_size; + NADim stride; + NADim dilation; + }; + + 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; + TMA_K tma_load_k; + TMA_V tma_load_v; + NADim qkv_shape; + NADim q_shape; + NADim kv_shape; + cute::tuple + na_params; // win, win_left, win_right, stride + bool is_fully_block_sparse; + bool has_kv_padding; + NADim dilation; + bool requires_qkv_fixup; + bool is_dilated; + int num_dilation_groups; + }; + + template + static Params to_underlying_arguments( + ProblemShape const& problem_shape, + Arguments const& args, + void* workspace) { + auto ptr_Q = args.ptr_Q; + auto ptr_K = args.ptr_K; + auto ptr_V = args.ptr_V; + auto dQ = args.dQ; + auto dK = args.dK; + auto dV = args.dV; + auto problem_shape_qk = problem_shape; + + auto params_qk = CollectiveMmaQK::to_underlying_arguments( + problem_shape_qk, + typename CollectiveMmaQK::Arguments{ + ptr_Q, + dQ, + ptr_K, + dK, + }, + /*workspace=*/nullptr); + + auto problem_shape_pv = select<0, 2, 1, 3>(problem_shape_qk); + auto params_pv = CollectiveMmaPV::to_underlying_arguments( + problem_shape_pv, + typename CollectiveMmaPV::Arguments{ + ptr_K, + dK, // never used, dummy + ptr_V, + select<1, 0, 2>(dV), + }, + /*workspace=*/nullptr); + + auto window_left = get_window_left(args.window_size); + auto window_right = get_window_right(args.window_size); + + bool requires_qkv_fixup = not evenly_divides(args.qkv_shape, args.dilation); + + return Params{ + params_qk.tma_load_a, + params_qk.tma_load_b, + params_pv.tma_load_b, + args.qkv_shape, + args.q_shape, + args.kv_shape, + make_tuple(args.window_size, window_left, window_right, args.stride), + fully_block_sparse( + args.qkv_shape, + args.window_size, + args.stride, + QTileShape{}, + KVTileShape{}), + /* has_kv_padding */ not evenly_divides(args.qkv_shape, KVTileShape{}), + args.dilation, + requires_qkv_fixup, + is_dilated(args.dilation), + size(args.dilation)}; + } + + 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_in, + 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) { + BlkCoord blk_coord_q = blk_coord_in; + BlkCoord blk_coord_kv = blk_coord_in; + + auto qkv_shape = params.qkv_shape; + if (params.requires_qkv_fixup) { + qkv_shape = Mask{}.correct_qkv_shape( + problem_shape, + params.qkv_shape, + blk_coord_in, + params.dilation, + params.num_dilation_groups); + } else if (params.is_dilated) { + qkv_shape = ceil_div(params.qkv_shape, params.dilation); + } + + auto [kv_start, num_tiles] = Mask{}.get_trip_count( + blk_coord_in, + MultiDimTileShape{}, + params.q_shape, + qkv_shape, + params.na_params); + + int mask_tile_count = size(num_tiles); + + auto kv_start_tile = ceil_div(kv_start, KVTileShape{}); + + auto kv_tiled = ceil_div(params.kv_shape, KVTileShape{}); + auto ctr = make_identity_tensor(num_tiles); + auto ctr_offset = domain_offset(kv_start_tile, ctr); + + auto kv_tiled_layout = make_layout(kv_tiled); + + auto iter_to_tile_map = [&ctr_offset, &kv_tiled_layout](int iter) { + return crd2idx(ctr_offset(iter), kv_tiled_layout); + }; + + 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) + + // compute gQ, sQ + // we load 2*get<0>(blk_coord), and 2*get<0>(blk_coord) + 1 + ThrMMA mma_qk = typename CollectiveMmaQK::TiledMma{}.get_slice(0); + Tensor mQ_qdl_p = + params.tma_load_q.get_tma_tensor(select<0, 2, 3>(problem_shape)); + + int q_offs_0 = 0; + + Tensor mQ_qdl = domain_offset( + make_coord(q_offs_0, _0{}, make_coord(_0{}, _0{})), mQ_qdl_p); + + Tensor gQ_qdl = local_tile( + mQ_qdl, TileShapeQK{}, make_coord(_, _, _), Step<_1, X, _1>{}); + Tensor tSgQ_qdl = mma_qk.partition_A(gQ_qdl); + Tensor sQ = + make_tensor(make_smem_ptr(storage.smem_q.data()), SmemLayoutQ{}); + auto [tQgQ_qdl, tQsQ] = tma_partition( + params.tma_load_q, + _0{}, + make_layout(_1{}), + group_modes<0, 3>(sQ), + group_modes<0, 3>(tSgQ_qdl)); + Tensor tQgQ = tQgQ_qdl(_, _, _0{}, get<2>(blk_coord_q)); + + // compute gK, sK + Tensor mK_kdl_p = + params.tma_load_k.get_tma_tensor(select<1, 2, 3>(problem_shape)); + + int kv_offs_0 = 0; + + Tensor mK_kdl = domain_offset( + make_coord(kv_offs_0, _0{}, make_coord(_0{}, _0{})), mK_kdl_p); + + Tensor gK_kdl = local_tile( + mK_kdl, TileShapeQK{}, make_coord(_, _, _), Step{}); + Tensor tSgK_kdl = mma_qk.partition_B(gK_kdl); + Tensor sK = + make_tensor(make_smem_ptr(storage.smem_k.data()), SmemLayoutK{}); + auto [tKgK_kdl, tKsK] = tma_partition( + params.tma_load_k, + _0{}, + make_layout(_1{}), + group_modes<0, 3>(sK), + group_modes<0, 3>(tSgK_kdl)); + Tensor tKgK = tKgK_kdl(_, _, _0{}, get<2>(blk_coord_kv)); + + // compute gV, sV + ThrMMA mma_pv = typename CollectiveMmaPV::TiledMma{}.get_slice(0); + Tensor mV_dkl_p = + params.tma_load_v.get_tma_tensor(select<2, 1, 3>(problem_shape)); + + Tensor mV_dkl = domain_offset( + make_coord(_0{}, kv_offs_0, make_coord(_0{}, _0{})), mV_dkl_p); + + Tensor gV_dkl = local_tile( + mV_dkl, TileShapePV{}, make_coord(_, _, _), Step{}); + Tensor tOgV_dkl = mma_pv.partition_B(gV_dkl); + Tensor sV = + make_tensor(make_smem_ptr(storage.smem_v.data()), SmemLayoutV{}); + auto [tVgV_dkl, tVsV] = tma_partition( + params.tma_load_v, + _0{}, + make_layout(_1{}), + group_modes<0, 3>(sV), + group_modes<0, 3>(tOgV_dkl)); + auto tVgV = tVgV_dkl(_, _0{}, _, get<2>(blk_coord_kv)); + + // 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_q); + int q1_index = 2 * get<0>(blk_coord_q) + 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_kv.producer_acquire(pipeline_kv_producer_state); + if (lane_predicate) { + auto tma_barrier = + pipeline_kv.producer_get_barrier(pipeline_kv_producer_state); + copy( + params.tma_load_k.with(*tma_barrier, 0), + tKgK(_, iter_to_tile_map(k_index)), + tKsK(_, pipeline_kv_producer_state.index())); + } + ++pipeline_kv_producer_state; + + // 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 + pipeline_kv.producer_acquire(pipeline_kv_producer_state); + if (lane_predicate) { + auto tma_barrier = + pipeline_kv.producer_get_barrier(pipeline_kv_producer_state); + copy( + params.tma_load_v.with(*tma_barrier, 0), + tVgV(_, iter_to_tile_map(k_index)), + tVsV(_, pipeline_kv_producer_state.index())); + } + ++pipeline_kv_producer_state; + k_index += 1; + + // loop: + mask_tile_count -= 1; + for (; mask_tile_count > 0; mask_tile_count -= 1) { + // Ki + pipeline_kv.producer_acquire(pipeline_kv_producer_state); + if (lane_predicate) { + auto tma_barrier = + pipeline_kv.producer_get_barrier(pipeline_kv_producer_state); + copy( + params.tma_load_k.with(*tma_barrier, 0), + tKgK(_, iter_to_tile_map(k_index)), + tKsK(_, pipeline_kv_producer_state.index())); + } + ++pipeline_kv_producer_state; + + // Vi + pipeline_kv.producer_acquire(pipeline_kv_producer_state); + if (lane_predicate) { + auto tma_barrier = + pipeline_kv.producer_get_barrier(pipeline_kv_producer_state); + copy( + params.tma_load_v.with(*tma_barrier, 0), + tVgV(_, iter_to_tile_map(k_index)), + tVsV(_, pipeline_kv_producer_state.index())); + } + ++pipeline_kv_producer_state; + k_index += 1; + } + } +}; + +} // namespace cutlass::fna::collective diff --git a/natten/csrc/include/natten/cuda/fna_blackwell/common/pow_2.hpp b/natten/csrc/include/natten/cuda/fna_blackwell/common/pow_2.hpp new file mode 100644 index 00000000..4964bb4c --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna_blackwell/common/pow_2.hpp @@ -0,0 +1,92 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2026 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::fna { + +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::fna + +namespace cute { + +template <> +struct is_integral : true_type {}; + +} // end namespace cute diff --git a/natten/csrc/include/natten/cuda/fna_blackwell/device/fna_bwd_sm100.hpp b/natten/csrc/include/natten/cuda/fna_blackwell/device/fna_bwd_sm100.hpp new file mode 100644 index 00000000..45b42153 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna_blackwell/device/fna_bwd_sm100.hpp @@ -0,0 +1,433 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 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 +#include "cute/tensor.hpp" +#include "cutlass/cutlass.h" +#include "cutlass/kernel_hardware_info.hpp" + +#include "natten/cuda/fna_blackwell/device/fna_sm100.hpp" +#include "natten/cuda/fna_blackwell/kernel/sm100_fna_bwd_kernel_tma_warpspecialized.hpp" + +// These kernels are identical; reuse from FMHA +#include "natten/cuda/fmha_blackwell/kernel/fmha_kernel_bwd_convert.hpp" +#include "natten/cuda/fmha_blackwell/kernel/fmha_kernel_bwd_sum_OdO.hpp" + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::fna::device { + +//////////////////////////////////////////////////////////////////////////////// +////////////////////////////// CUTLASS 3.x API ///////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +template < + class ProblemShape, + class Element, + class ElementAccumulator, + class TileShape, + class Mask, + class QTileShape, + class KVTileShape, + class NADim> +class FnaBwdSm100 { + public: + /// Argument structure: User API + struct Arguments { + // Q K D D_VO HB + ProblemShape problem_shape; + + const Element* ptr_Q; + cute::tuple, int64_t>> + stride_Q; + const Element* ptr_K; + cute::tuple, int64_t>> + stride_K; + const Element* ptr_V; + cute::tuple, int64_t>> + stride_V; + + const Element* ptr_O; + cute::tuple, int64_t>> + stride_O; + const ElementAccumulator* ptr_LSE; + // NATTEN has a different LSE layout + cute::tuple, int64_t>> stride_LSE; + + const Element* ptr_dO; + cute::tuple, int64_t>> + stride_dO; + + Element* ptr_dQ; + cute::tuple, int64_t>> + stride_dQ; + Element* ptr_dK; + cute::tuple, int64_t>> + stride_dK; + Element* ptr_dV; + cute::tuple, int64_t>> + stride_dV; + + // (F)NA parameters + NADim q_shape; + NADim kv_shape; + NADim qkv_shape; + NADim window_size; + NADim stride; + NADim dilation; + + ElementAccumulator softmax_scale; + + cutlass::KernelHardwareInfo hw_info; + }; + + using OperationSumOdO = cutlass::fna::device::FnaSm100< + cutlass::fmha::kernel:: + FmhaKernelBwdSumOdO>; + using OperationConvert = cutlass::fna::device::FnaSm100< + cutlass::fmha::kernel:: + FmhaKernelBwdConvert>; + + using Operation = cutlass::fna::device::FnaSm100< + cutlass::fna::kernel::Sm100FnaBwdKernelTmaWarpSpecialized< + ProblemShape, + Element, + ElementAccumulator, + TileShape, + Mask, + QTileShape, + KVTileShape, + NADim>>; + using Kernel = typename Operation::Kernel; + + struct Params { + OperationSumOdO op_sum_OdO; + Operation op; + OperationConvert op_convert; + ElementAccumulator* dQ_acc; + size_t dQ_acc_size; + }; + + private: + Params params_; + + static typename OperationSumOdO::Arguments to_sum_OdO_arguments( + Arguments const& args, + ElementAccumulator* sum_odo = nullptr, + ElementAccumulator* scaled_lse = nullptr) { + using namespace cute; + auto [Q_, K, D, D_VO, HB] = args.problem_shape; + auto [H, B] = HB; + auto [H_R, H_K] = H; + D = cutlass::round_up(D, 8); // Alignment + int Q = cutlass::round_up(static_cast(Q_), 8); // Alignment + auto stride_sum_OdO = make_stride( + _1{}, + make_stride( + make_stride(Q, Q * H_R), + B == 1 + ? 0 + : static_cast(Q) * static_cast(H_R * H_K))); + auto stride_scaled_lse = make_stride( + _1{}, + make_stride( + make_stride(Q, Q * H_R), + B == 1 + ? 0 + : static_cast(Q) * static_cast(H_R * H_K))); + auto log2_e = log2f(expf(1.0f)); + return typename OperationSumOdO::Arguments{ + args.problem_shape, + args.ptr_O, + args.stride_O, + args.ptr_dO, + args.stride_dO, + sum_odo, + stride_sum_OdO, + args.ptr_LSE, + args.stride_LSE, + scaled_lse, + stride_scaled_lse, + -1.0f, + -log2_e}; + } + + static typename OperationConvert::Arguments to_convert_arguments( + Arguments const& args, + ElementAccumulator* src = nullptr) { + using namespace cute; + auto [Q_, K, D, D_VO, HB] = args.problem_shape; + auto [H, B] = HB; + auto [H_R, H_K] = H; + D = cutlass::round_up(D, 8); // Alignment + int Q = cutlass::round_up(static_cast(Q_), 8); // Alignment + auto stride_src_dQ = make_stride( + D, + _1{}, + make_stride( + make_stride(D * Q, D * Q * H_R), + B == 1 ? 0 + : static_cast(D * H_R * H_K) * + static_cast(Q))); + return typename OperationConvert::Arguments{ + args.problem_shape, + src, + stride_src_dQ, + nullptr, + args.stride_dK, + nullptr, + args.stride_dV, + args.ptr_dQ, + args.stride_dQ, + nullptr, + args.stride_dK, + nullptr, + args.stride_dV, + args.softmax_scale}; + } + + static typename Operation::Arguments to_bwd_arguments( + Arguments const& args, + ElementAccumulator* sum_OdO = nullptr, + cute::tuple, int64_t>> const& + stride_sum_OdO = {}, + ElementAccumulator* scaled_lse = nullptr, + cute::tuple, int64_t>> const& + stride_scaled_lse = {}, + ElementAccumulator* dQ_acc = nullptr, + cute::tuple< + int, + cute::_1, + cute::tuple, int64_t>> const& stride_dQ = {}) { + return typename Operation::Arguments{ + args.problem_shape, + // FNA args + {args.q_shape, + args.kv_shape, + args.qkv_shape, + args.window_size, + args.stride, + args.dilation}, + {args.ptr_Q, + args.stride_Q, + args.ptr_K, + args.stride_K, + args.ptr_V, + args.stride_V, + args.ptr_dO, + args.stride_dO, + scaled_lse, + stride_scaled_lse, + sum_OdO, + stride_sum_OdO, + dQ_acc, + stride_dQ, + args.softmax_scale}, + {args.ptr_dK, args.stride_dK, args.ptr_dV, args.stride_dV}, + args.hw_info}; + } + + public: + /// Determines whether the GEMM can execute the given problem. + static Status can_implement(Arguments const& args) { + Status status = Status::kSuccess; + + status = OperationSumOdO::can_implement(to_sum_OdO_arguments(args)); + if (status != Status::kSuccess) { + return status; + } + + status = OperationConvert::can_implement(to_convert_arguments(args)); + if (status != Status::kSuccess) { + return status; + } + + status = Operation::can_implement(to_bwd_arguments(args)); + if (status != Status::kSuccess) { + return status; + } + + return status; + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const& args) { + auto [Q_, K, D_, D_VO, HB] = args.problem_shape; + auto [H_, B_] = product_each(HB); + size_t B = static_cast(B_); + size_t H = static_cast(H_); + size_t D = cutlass::round_up(static_cast(D_), 8); // Alignment + size_t Q = cutlass::round_up(static_cast(Q_), 8); // Alignment + size_t workspace_bytes = 0; + // OdO vector + workspace_bytes += B * H * Q * sizeof(ElementAccumulator); + // scaled LSE vector + workspace_bytes += B * H * Q * sizeof(ElementAccumulator); + // FP32 versions of outputs that are churned (start off with Q only) + workspace_bytes += B * H * Q * D * sizeof(ElementAccumulator); + return workspace_bytes; + } + + /// Initializes state from arguments. + Status initialize_split( + Arguments const& args, + void* workspace_dQ, + void* workspace_sum_OdO, + void* workspace_scaled_lse, + cudaStream_t stream = nullptr) { + CUTLASS_TRACE_HOST( + "FnaDeviceBwd::initialize_split() - workspace_dQ=" + << workspace_dQ << ", workspace_sum_OdO=" << workspace_sum_OdO + << "stream: " << (stream ? "non-null" : "null")); + + auto [Q_, K, D_, D_VO, HB] = args.problem_shape; + auto [H_, B_] = product_each(HB); + size_t B = static_cast(B_); + size_t H = static_cast(H_); + size_t D = cutlass::round_up(static_cast(D_), 8); // Alignment + size_t Q = cutlass::round_up(static_cast(Q_), 8); // Alignment + ElementAccumulator* sum_OdO = + reinterpret_cast(workspace_sum_OdO); + ElementAccumulator* scaled_lse = + reinterpret_cast(workspace_scaled_lse); + ElementAccumulator* dQ_acc = + reinterpret_cast(workspace_dQ); + params_.dQ_acc = dQ_acc; + params_.dQ_acc_size = B * H * Q * D * sizeof(ElementAccumulator); + auto args_sum_OdO = to_sum_OdO_arguments(args, sum_OdO, scaled_lse); + auto args_convert = to_convert_arguments(args, dQ_acc); + params_.op_sum_OdO.initialize(args_sum_OdO, nullptr, stream); + params_.op_convert.initialize(args_convert, nullptr, stream); + auto args_bwd = to_bwd_arguments( + args, + sum_OdO, + args_sum_OdO.stride_sum_OdO, + scaled_lse, + args_sum_OdO.stride_scaled_lse, + dQ_acc, + args_convert.stride_src_dQ); + params_.op.initialize(args_bwd, nullptr, stream); + + return Status::kSuccess; + } + + /// Initializes state from arguments. + Status initialize( + Arguments const& args, + void* workspace = nullptr, + cudaStream_t stream = nullptr) { + CUTLASS_TRACE_HOST( + "FnaDeviceBwd::initialize() - workspace " + << workspace << ", stream: " << (stream ? "non-null" : "null")); + + auto [Q_, K, D_, D_VO, HB] = args.problem_shape; + auto [H_, B_] = product_each(HB); + size_t B = static_cast(B_); + size_t H = static_cast(H_); + size_t D = cutlass::round_up(static_cast(D_), 8); // Alignment + size_t Q = cutlass::round_up(static_cast(Q_), 8); // Alignment + char* workspace_chr = reinterpret_cast(workspace); + ElementAccumulator* sum_OdO = + reinterpret_cast(workspace_chr); + workspace_chr += B * H * Q * sizeof(ElementAccumulator); + ElementAccumulator* scaled_lse = + reinterpret_cast(workspace_chr); + workspace_chr += B * H * Q * sizeof(ElementAccumulator); + ElementAccumulator* dQ_acc = + reinterpret_cast(workspace_chr); + return initialize_split(args, dQ_acc, sum_OdO, scaled_lse, stream); + } + + /// 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("FnaDeviceBwd::run()"); + + Status result = Status::kSuccess; + result = params.op_sum_OdO.run(stream); + if (result != Status::kSuccess) { + return result; + } + + auto cuda_result = + cudaMemsetAsync(params.dQ_acc, 0, params.dQ_acc_size, stream); + if (cuda_result != cudaSuccess) { + return Status::kErrorInternal; + } + + result = params.op.run(stream); + if (result != Status::kSuccess) { + return result; + } + + result = params.op_convert.run(stream); + if (result != Status::kSuccess) { + return result; + } + + 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; + } + + /// 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); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::fna::device + +//////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fna_blackwell/device/fna_sm100.hpp b/natten/csrc/include/natten/cuda/fna_blackwell/device/fna_sm100.hpp new file mode 100644 index 00000000..2e466871 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna_blackwell/device/fna_sm100.hpp @@ -0,0 +1,286 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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 + +#include +// common +#include "cutlass/cutlass.h" + +#if !defined(__CUDACC_RTC__) +#include "cutlass/cluster_launch.hpp" +#include "cutlass/trace.h" +#endif // !defined(__CUDACC_RTC__) + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::fna::device { + +//////////////////////////////////////////////////////////////////////////////// +////////////////////////////// CUTLASS 3.x API ///////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +template +class FnaSm100 { + 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("FNA::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_sm100, + 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_sm100, + 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( + "FNA::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_sm100, + 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("FNA()::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) { + CUTLASS_TRACE_HOST("FNA::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_sm100; + void* kernel_params[] = {¶ms}; + launch_result = ClusterLauncher::launch( + grid, cluster, block, smem_size, stream, kernel, kernel_params); + } else { + launch_result = Status::kSuccess; + device_kernel_sm100<<>>(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::fna::device + +//////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fna_blackwell/fna_backward.cuh b/natten/csrc/include/natten/cuda/fna_blackwell/fna_backward.cuh new file mode 100644 index 00000000..40b748ca --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna_blackwell/fna_backward.cuh @@ -0,0 +1,249 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ + +#pragma once + +#include "cute/tensor.hpp" + +#include "cutlass/cutlass.h" +#include "cutlass/kernel_hardware_info.h" + +#include "natten/cuda/fna_blackwell/collective/fna_fusion_bwd.hpp" +#include "natten/cuda/fna_blackwell/device/fna_bwd_sm100.hpp" + +namespace natten { +namespace cuda { +namespace fna_blackwell { + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +using namespace cute; +using namespace cutlass::fna; + +template < + typename Element, + class Causal, + class QTileShape, + class KVTileShape, + class TileShape> +struct KernelBackward { + static_assert( + rank(QTileShape{}) == 1 || rank(QTileShape{}) == 2 || + rank(QTileShape{}) == 3); + static_assert(rank(QTileShape{}) == rank(KVTileShape{})); + static_assert(rank(QTileShape{}) == rank(Causal{})); + using NADim = std::conditional_t< + rank(QTileShape{}) == 3, + cute::tuple, + std::conditional_t< + rank(QTileShape{}) == 2, + cute::tuple, + cute::tuple>>; + + using ElementAccumulator = float; + + // Q K D D_VO ((H_R, H_K) B) + using ProblemShapeType = + cute::tuple, int>>; + + using Operation = cutlass::fna::device::FnaBwdSm100< + ProblemShapeType, + Element, + ElementAccumulator, + TileShape, + collective::NeighborhoodAttentionBackwardMask, + QTileShape, + KVTileShape, + NADim>; + + using Arguments = typename Operation::Arguments; + + Operation op; + + Arguments initialize( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch, + int seqlen_Q, + int seqlen_KV, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + NADim q_shape, + NADim kv_shape, + NADim qkv_shape, + NADim window_size, + NADim stride, + NADim dilation, + // init/launch params + int device_id) { + auto problem_shape_regular = cute::make_tuple( + seqlen_Q, + seqlen_KV, + // head dims are checked to be multiples of 8 for fp16/bf16 and 16 for + // fp8, to meet the 128-bit alignment constraint. No action needed here. + dim, + dim, // dim_value -- if different from dim, needs the MLA kernel + cute::make_tuple(make_tuple(heads_q / heads_kv, heads_kv), batch)); + + ProblemShapeType problem_shape_launch; + decltype(problem_shape_regular) problem_shape_memory; + + problem_shape_memory = problem_shape_regular; + problem_shape_launch = problem_shape_regular; + int SQ = size<0>(problem_shape_memory); + int SK = size<1>(problem_shape_memory); + int D = size<2>(problem_shape_memory); + int D_VO = size<3>(problem_shape_memory); + auto HB = get<4, 0>(problem_shape_memory); + auto [H_R, H_K] = HB; + int B = size<4, 1>(problem_shape_memory); + + // heads last profile, with torch's "contiguous layout" + // shape: (batch, seqlen, heads, dim) + // stride: (dim*heads*seqlen, dim*heads, dim, 1) + auto stride_Q = make_stride( + H_R * H_K * D, + _1{}, + make_stride( + make_stride(D, D * H_R), + B == 1 ? 0L + : static_cast(D * H_R * H_K) * + static_cast(SQ))); + auto stride_K = make_stride( + H_K * D, + _1{}, + make_stride( + make_stride(_0{}, D), + B == 1 ? 0L + : static_cast(D * H_K) * static_cast(SK))); + auto stride_V = make_stride( + H_K * D_VO, + _1{}, + make_stride( + make_stride(_0{}, D_VO), + B == 1 + ? 0L + : static_cast(D_VO * H_K) * static_cast(SK))); + auto stride_O = make_stride( + H_R * H_K * D_VO, + _1{}, + make_stride( + make_stride(D_VO, D_VO * H_R), + B == 1 ? 0L + : static_cast(D_VO * H_R * H_K) * + static_cast(SQ))); + auto stride_LSE = make_stride( + H_K * H_R, + make_stride( + make_stride(_1{}, H_R), + B == 1 + ? 0L + : static_cast(SQ) * static_cast(H_R * H_K))); + + cutlass::KernelHardwareInfo hw_info; + hw_info.device_id = device_id; + hw_info.sm_count = + cutlass::KernelHardwareInfo::query_device_multiprocessor_count( + hw_info.device_id); + + Arguments arguments{ + problem_shape_launch, + reinterpret_cast(ptr_Q), + stride_Q, + reinterpret_cast(ptr_K), + stride_K, + reinterpret_cast(ptr_V), + stride_V, + reinterpret_cast(ptr_O), + stride_O, + reinterpret_cast(ptr_LSE), + stride_LSE, + reinterpret_cast(ptr_dO), + stride_O, + reinterpret_cast(ptr_dQ), + stride_Q, + reinterpret_cast(ptr_dK), + stride_K, + reinterpret_cast(ptr_dV), + stride_V, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + attn_scale, + hw_info}; + + return arguments; + } + + size_t get_workspace_size(Arguments const& arguments) { + return Operation::get_workspace_size(arguments); + } + + void run( + Arguments const& arguments, + void* workspace_ptr, + cudaStream_t stream) { + 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; + return; + } + + status = op.initialize(arguments, workspace_ptr, stream); + if (status != cutlass::Status::kSuccess) { + std::cerr + << "Failed to initialize the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(cudaGetLastError()) << std::endl; + return; + } + + // 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; + } + } +}; + +#endif // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten diff --git a/natten/csrc/include/natten/cuda/fna_blackwell/fna_forward.cuh b/natten/csrc/include/natten/cuda/fna_blackwell/fna_forward.cuh new file mode 100644 index 00000000..20838538 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna_blackwell/fna_forward.cuh @@ -0,0 +1,270 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ + +#pragma once + +#include "cute/tensor.hpp" + +#include "cutlass/cutlass.h" +#include "cutlass/kernel_hardware_info.h" + +#include "natten/cuda/fna_blackwell/collective/fna_fusion.hpp" +#include "natten/cuda/fna_blackwell/collective/sm100_fna_fwd_epilogue_tma_warpspecialized.hpp" +#include "natten/cuda/fna_blackwell/collective/sm100_fna_fwd_mainloop_tma_warpspecialized.hpp" +#include "natten/cuda/fna_blackwell/device/fna_sm100.hpp" +#include "natten/cuda/fna_blackwell/kernel/sm100_fna_fwd_kernel_tma_warpspecialized.hpp" + +#include "natten/cuda/fmha_blackwell/kernel/fmha_tile_scheduler.hpp" + +namespace natten { +namespace cuda { +namespace fna_blackwell { + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +using namespace cute; +using namespace cutlass::fna::kernel; +using namespace cutlass::fna::collective; +using namespace cutlass::fna; + +template < + typename Element, + class Causal, + class QTileShape, + class KVTileShape, + class TileShape, + bool kIsPersistent> +struct KernelForward { + static_assert( + rank(QTileShape{}) == 1 || rank(QTileShape{}) == 2 || + rank(QTileShape{}) == 3); + static_assert(rank(QTileShape{}) == rank(KVTileShape{})); + static_assert(rank(QTileShape{}) == rank(Causal{})); + using NADim = std::conditional_t< + rank(QTileShape{}) == 3, + cute::tuple, + std::conditional_t< + rank(QTileShape{}) == 2, + cute::tuple, + cute::tuple>>; + + using ElementAccumulatorQK = float; + using ElementAccumulatorPV = float; + using ElementOut = Element; + + // Q K D (B H) + using ProblemShapeType = + cute::tuple, int>>; + + using StrideQ = cute:: + tuple, int64_t>>; // Q D (H_G + // H_R B) + using StrideK = cute:: + tuple, int64_t>>; // K D (H_G + // H_R B) + using StrideV = StrideK; + using StrideO = StrideQ; + using StrideLSE = + cute::tuple, int64_t>>; // Q (H_G + // H_R B) + + using TileScheduler = std::conditional_t< + kIsPersistent, + cutlass::fmha::kernel::PersistentTileScheduler, + cutlass::fmha::kernel::IndividualTileScheduler>; + + using Mainloop = + cutlass::fna::collective::Sm100FnaFwdMainloopTmaWarpspecialized< + Element, + ElementAccumulatorQK, + ElementAccumulatorPV, + TileShape, + StrideQ, + StrideK, + StrideV, + NeighborhoodAttentionMask, + QTileShape, + KVTileShape, + NADim>; + using Operation = cutlass::fna::device::FnaSm100< + cutlass::fna::kernel::Sm100FnaFwdKernelTmaWarpspecialized< + ProblemShapeType, + Mainloop, + cutlass::fna::collective::Sm100FnaFwdEpilogueTmaWarpspecialized< + ElementOut, + ElementAccumulatorPV, + typename Mainloop::TileShapePV, + StrideO, + StrideLSE>, + TileScheduler>>; + + using Arguments = typename Operation::Arguments; + + Operation op; + + Arguments initialize( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch, + int seqlen_Q, + int seqlen_KV, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + NADim q_shape, + NADim kv_shape, + NADim qkv_shape, + NADim window_size, + NADim stride, + NADim dilation, + // init/launch params + int device_id) { + auto problem_shape_regular = cute::make_tuple( + seqlen_Q, + seqlen_KV, + // head dims are checked to be multiples of 8 for fp16/bf16 and 16 for + // fp8, to meet the 128-bit alignment constraint. No action needed here. + dim, + cute::make_tuple( + cute::make_tuple(heads_q / heads_kv, heads_kv), batch)); + + ProblemShapeType problem_shape_launch; + decltype(problem_shape_regular) problem_shape_memory; + + problem_shape_memory = problem_shape_regular; + problem_shape_launch = problem_shape_regular; + + int SQ = size<0>(problem_shape_memory); + int SK = size<1>(problem_shape_memory); + int D = size<2>(problem_shape_memory); + int H = size<3, 0>(problem_shape_memory); + int H_K = size<3, 0, 1>(problem_shape_memory); + int H_Q = size<3, 0, 0>(problem_shape_memory); + int B = size<3, 1>(problem_shape_memory); + + // heads last profile, with torch's "contiguous layout" + // shape: (batch, seqlen, heads, dim) + // stride: (dim*heads*seqlen, dim*heads, dim, 1) + auto stride_Q = make_stride( + H * D, + _1{}, + make_stride( + make_stride(D, H_Q * D), + static_cast(H * D) * static_cast(SQ))); + auto stride_O = stride_Q; + auto stride_K = make_stride( + H_K * D, + _1{}, + make_stride( + make_stride(_0{}, D), + static_cast(H_K * D) * static_cast(SK))); + auto stride_V = stride_K; + auto stride_LSE = make_stride( + H, + make_stride( + make_stride(_1{}, H_Q), + static_cast(SQ) * static_cast(H))); + + cutlass::KernelHardwareInfo hw_info; + hw_info.device_id = device_id; + hw_info.sm_count = + cutlass::KernelHardwareInfo::query_device_multiprocessor_count( + hw_info.device_id); + + Arguments arguments{ + problem_shape_launch, + {reinterpret_cast(ptr_Q), + stride_Q, + reinterpret_cast(ptr_K), + stride_K, + reinterpret_cast(ptr_V), + stride_V, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + attn_scale}, + {reinterpret_cast(ptr_O), + stride_O, + reinterpret_cast(ptr_LSE), + stride_LSE}, + hw_info}; + + return arguments; + } + + size_t get_workspace_size(Arguments const& arguments) { + return Operation::get_workspace_size(arguments); + } + + void run( + Arguments const& arguments, + void* workspace_ptr, + cudaStream_t stream) { + 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; + return; + } + + status = op.initialize(arguments, workspace_ptr, stream); + if (status != cutlass::Status::kSuccess) { + std::cerr + << "Failed to initialize the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(cudaGetLastError()) << std::endl; + return; + } + + // 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; + } + +#if 0 + cudaError_t result = cudaDeviceSynchronize(); + if (result != cudaSuccess) { + std::cerr << "Error running the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(result) << std::endl; + return; + } +#endif + } +}; + +#endif // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +} // namespace fna_blackwell +} // namespace cuda +} // namespace natten diff --git a/natten/csrc/include/natten/cuda/fna_blackwell/kernel/sm100_fna_bwd_kernel_tma_warpspecialized.hpp b/natten/csrc/include/natten/cuda/fna_blackwell/kernel/sm100_fna_bwd_kernel_tma_warpspecialized.hpp new file mode 100644 index 00000000..bd8d19c0 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna_blackwell/kernel/sm100_fna_bwd_kernel_tma_warpspecialized.hpp @@ -0,0 +1,2389 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 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 "cute/arch/simd_sm100.hpp" +#include "cute/tensor.hpp" + +#include "cutlass/arch/arch.h" +#include "cutlass/arch/memory_sm80.h" +#include "cutlass/gemm/collective/collective_builder.hpp" + +#include "natten/cuda/fna_blackwell/collective/fna_common.hpp" +#include "natten/cuda/fna_blackwell/collective/fna_fusion.hpp" +#include "natten/cuda/fna_blackwell/collective/fna_fusion_bwd.hpp" + +namespace cutlass::fna::kernel { + +using namespace cutlass::fna::collective; + +using namespace cute; + +template < + class ProblemShape, + class Element, + class ElementAcc, + class TileShape, + class Mask, + class QTileShape, + class KVTileShape, + class NADim> +struct Sm100FnaBwdKernelTmaWarpSpecialized { + using MultiDimTileShape = cute::tuple; + + static_assert( + size(QTileShape{}) == get<0>(TileShape{}), + "QTileShape doesn't match the size of Q tile in the FMHA kernel."); + static_assert( + size(KVTileShape{}) == get<1>(TileShape{}), + "KVTileShape doesn't match the size of KV tile in the FMHA kernel."); + + using TileShapeQ = decltype(get<0>(TileShape{})); + static_assert(std::is_same_v, "tile shape K must be 128"); + using TileShapeK = decltype(get<1>(TileShape{})); + static_assert(std::is_same_v, "tile shape K must be 128"); + using TileShapeDQK = decltype(get<2>(TileShape{})); + using TileShapeDVO = decltype(get<2>(TileShape{})); + + using TmemAllocator = cute::TMEM::Allocator1Sm; + struct TmemAllocation { + static constexpr uint32_t kDK = 0; // TileShapeK x TileShapeDQK x acc + static constexpr uint32_t kDV = + kDK + TileShapeDQK{}; // TileShapeK x TileShapeDVO x acc + static constexpr uint32_t kDQ = + kDV + TileShapeDVO{}; // TileShapeQ x TileShapeDQK x acc + static constexpr uint32_t kDP = kDQ; // TileShapeK x TileShapeQ x inp + static constexpr uint32_t kS = kDQ + max(TileShapeQ{}, TileShapeDQK{}); + static constexpr uint32_t kP = kS; + static constexpr uint32_t kTotal = kS + TileShapeQ{}; + }; + + static_assert( + static_cast(TmemAllocation::kTotal) <= + TmemAllocator::Sm100TmemCapacityColumns, + "using too much tmem"); + + enum class WarpRole { + Empty = 0x0, + Load = 0x1, + Mma = 0x2, + Compute = 0x3, + Reduce = 0x4 + }; + + static constexpr unsigned long long kWarpAssignment = 0x12'3333'3333'4444ull; + static constexpr int kNumComputeWarps = 8; + static constexpr int kNumReduceWarps = 4; + CUTLASS_DEVICE WarpRole warp_idx_to_role(int warp_idx) { + return static_cast((kWarpAssignment >> (4 * warp_idx)) & 0xF); + } + + struct RegisterAllocation { + static constexpr int kWarpgroup0 = 160 - 8; + static constexpr int kWarpgroup1 = 128; + static constexpr int kWarpgroup2 = 96; + static constexpr int kReduce = kWarpgroup0; + static constexpr int kCompute = kWarpgroup1; + static constexpr int kMma = kWarpgroup2; + static constexpr int kEmpty = kWarpgroup2; + static constexpr int kLoad = kWarpgroup2; + + static_assert(kWarpgroup0 + 2 * kWarpgroup1 + kWarpgroup2 <= 512); + }; + + using ArchTag = cutlass::arch::Sm100; + + using ClusterShape = Shape<_1, _1, _1>; + using Schedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; + + static constexpr int MinBlocksPerMultiprocessor = 1; + static constexpr int kNumWarps = kNumComputeWarps + kNumReduceWarps + 4; + static constexpr int MaxThreadsPerBlock = NumThreadsPerWarp * kNumWarps; + + static constexpr int Alignment = 128 / sizeof_bits_v; + static constexpr int kStages = 2; + + using TensorStrideContiguousK = + Stride, int64_t>>; + using TensorStrideContiguousMN = + Stride<_1, int, Stride, int64_t>>; + using TensorStrideContiguousK_GQA = + Stride, int64_t>>; + using TensorStrideContiguousMN_GQA = + Stride<_1, int, Stride, int64_t>>; + + // compute S + using CollectiveMmaKQ = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, + cutlass::arch::OpClassTensorOp, + Element, + TensorStrideContiguousK_GQA, + Alignment, + Element, + TensorStrideContiguousK, + Alignment, + ElementAcc, + Shape, + ClusterShape, + cutlass::gemm::collective::StageCount, + Schedule>::CollectiveOp; + using TileShapeKQ = typename CollectiveMmaKQ::TileShape; + using TiledMmaKQ = typename CollectiveMmaKQ::TiledMma; + + // compute dP + using CollectiveMmaVDO = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, + cutlass::arch::OpClassTensorOp, + Element, + TensorStrideContiguousK_GQA, + Alignment, + Element, + TensorStrideContiguousK, + Alignment, + ElementAcc, + Shape, + ClusterShape, + cutlass::gemm::collective::StageCount, + Schedule>::CollectiveOp; + using TileShapeVDO = typename CollectiveMmaVDO::TileShape; + using TiledMmaVDO = typename CollectiveMmaVDO::TiledMma; + + // compute dV + using CollectiveMmaPDO = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, + cutlass::arch::OpClassTensorOp, + // needs to match ordering of S calculation + Element, + TensorStrideContiguousK, + Alignment, + Element, + TensorStrideContiguousMN, + Alignment, + ElementAcc, + Shape, + ClusterShape, + cutlass::gemm::collective::StageCount, + Schedule>::CollectiveOp; + using TileShapePDO = typename CollectiveMmaPDO::TileShape; + using TiledMmaPDO = + decltype(to_tiled_mma_sm100_ts(typename CollectiveMmaPDO::TiledMma{})); + + // compute dK + using CollectiveMmaDSQ = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, + cutlass::arch::OpClassTensorOp, + // somewhat arbitrary since we dump to smem, need to agree with the + // next one + Element, + TensorStrideContiguousK, + Alignment, + Element, + TensorStrideContiguousMN, + Alignment, + ElementAcc, + Shape, + ClusterShape, + cutlass::gemm::collective::StageCount, + Schedule>::CollectiveOp; + using TileShapeDSQ = typename CollectiveMmaDSQ::TileShape; + using TiledMmaDSQ = typename CollectiveMmaDSQ::TiledMma; + + // compute dQ + using CollectiveMmaDSK = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, + cutlass::arch::OpClassTensorOp, + // somewhat arbitrary since we dump to smem, need to agree with the + // previous one + Element, + TensorStrideContiguousMN, + Alignment, + Element, + TensorStrideContiguousMN_GQA, + Alignment, + ElementAcc, + Shape, + ClusterShape, + cutlass::gemm::collective::StageCount, + Schedule>::CollectiveOp; + using TileShapeDSK = typename CollectiveMmaDSK::TileShape; + using TiledMmaDSK = typename CollectiveMmaDSK::TiledMma; + + // pipelines are named Pipeline + static constexpr int kStagesComputeSmem = 1; + using PipelineLoadMmaQ = PipelineTmaUmmaAsync<2, ClusterShape>; + using PipelineLoadMmaDO = PipelineTmaUmmaAsync<1, ClusterShape>; + using PipelineLoadComputeLSE = PipelineAsync<1>; + using PipelineLoadComputeSumOdO = PipelineAsync<1>; + using PipelineMmaComputeS = PipelineUmmaAsync<1>; + using PipelineMmaComputeDP = PipelineUmmaAsync<1>; + using PipelineMmaReduceDQ = PipelineUmmaAsync<1>; + using PipelineComputeMmaP = PipelineUmmaConsumerAsync<1>; + using PipelineComputeMmaDS = PipelineUmmaConsumerAsync; + using PipelineMmaComputeDKDV = PipelineUmmaAsync<2>; + static constexpr int kStagesReduceTmaStore = 2; + using PipelineReduceTmaStore = PipelineTmaStore; + + struct PipelineStorage { + alignas(16) typename PipelineLoadMmaQ::SharedStorage load_mma_q; + alignas(16) typename PipelineLoadMmaDO::SharedStorage load_mma_do; + alignas(16) typename PipelineLoadComputeLSE::SharedStorage load_compute_lse; + alignas(16) + typename PipelineLoadComputeSumOdO::SharedStorage load_compute_sum_odo; + alignas(16) typename PipelineMmaComputeS::SharedStorage mma_compute_s; + alignas(16) typename PipelineMmaComputeDP::SharedStorage mma_compute_dp; + alignas(16) typename PipelineMmaReduceDQ::SharedStorage mma_reduce_dq; + alignas(16) typename PipelineComputeMmaP::SharedStorage compute_mma_p; + alignas(16) typename PipelineComputeMmaDS::SharedStorage compute_mma_ds; + alignas(16) typename PipelineMmaComputeDKDV::SharedStorage mma_compute_dkdv; + }; + + template + static CUTE_DEVICE constexpr auto restage( + Layout const& layout, + Stages stages = {}) { + return composition(layout, make_tuple(_, _, _, make_layout(stages))); + } + + using SmemLayoutK = + decltype(restage(typename CollectiveMmaKQ::SmemLayoutA{})); + using SmemLayoutV = + decltype(restage(typename CollectiveMmaVDO::SmemLayoutA{})); + using SmemLayoutQ = + decltype(restage(typename CollectiveMmaKQ::SmemLayoutB{}, _2{})); + using SmemLayoutDO = + decltype(restage(typename CollectiveMmaVDO::SmemLayoutB{}, _1{})); + using SmemLayoutDS = decltype(restage( + typename CollectiveMmaDSK::SmemLayoutA{}, + Int{})); + using SmemLayoutLSE = Layout>; + using SmemLayoutSumOdO = Layout>; + + using SmemLayoutQT = + decltype(restage(typename CollectiveMmaDSQ::SmemLayoutB{}, _2{})); + using SmemLayoutKT = + decltype(restage(typename CollectiveMmaDSK::SmemLayoutB{})); + using SmemLayoutDST = decltype(restage( + typename CollectiveMmaDSQ::SmemLayoutA{}, + Int{})); + using SmemLayoutDOT = + decltype(restage(typename CollectiveMmaPDO::SmemLayoutB{}, _1{})); + + using TileShapeDQ = _32; + using SmemAtomDQ = + decltype(cutlass::gemm::collective::detail::sm100_smem_selector< + cute::UMMA::Major::K, + ElementAcc, + TileShapeQ, + TileShapeDQ>()); + using SmemShapeDQ = + Shape>; + using SmemLayoutDQ = + decltype(tile_to_shape(SmemAtomDQ{}, SmemShapeDQ{}, Step<_2, _1, _3>{})); + + struct TensorStorage { + union { + alignas(2048) cute::array> smem_k; + alignas(2048) cute::array> smem_k_t; + }; + alignas(2048) cute::array> smem_v; + union { + alignas(2048) cute::array> smem_q; + alignas(2048) cute::array> smem_q_t; + }; + union { + alignas(2048) cute::array> smem_do; + alignas( + 2048) cute::array> smem_do_t; + }; + union { + alignas(2048) cute::array> smem_ds; + alignas( + 2048) cute::array> smem_ds_t; + }; + alignas(1024) cute::array> smem_dq; + alignas(16) cute::array> smem_lse; + alignas(16) + cute::array> smem_sum_odo; + }; + + static constexpr int kTransactionsBytesLoadQ = cutlass::bits_to_bytes( + cosize(take<0, 3>(SmemLayoutQ{})) * cute::sizeof_bits_v); + static constexpr int kTransactionsBytesLoadDO = cutlass::bits_to_bytes( + cosize(take<0, 3>(SmemLayoutDO{})) * cute::sizeof_bits_v); + + static constexpr int kTransactionsBytesLoadK = cutlass::bits_to_bytes( + cosize(take<0, 3>(SmemLayoutK{})) * cute::sizeof_bits_v); + static constexpr int kTransactionsBytesLoadV = cutlass::bits_to_bytes( + cosize(take<0, 3>(SmemLayoutV{})) * cute::sizeof_bits_v); + + struct SharedStorage { + TensorStorage tensors; + PipelineStorage pipelines; + uint32_t tmem_base_ptr; + }; + + // this is tight enough that it won't work with sizeof due to padding for + // alignment + static constexpr int SharedStorageSize = + offsetof(SharedStorage, tmem_base_ptr) + sizeof(uint32_t); + static_assert( + SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, + "using too much smem"); + + using TensorStride = TensorStrideContiguousK; // S D (H B) + using TensorStride_GQA = TensorStrideContiguousK_GQA; + using RowTensorStride = + Stride<_1, Stride, int64_t>>; // S (H B) + + struct MainloopArguments { + const Element* ptr_q; + TensorStride stride_q; + const Element* ptr_k; + TensorStride_GQA stride_k; + const Element* ptr_v; + TensorStride_GQA stride_v; + const Element* ptr_do; + TensorStride stride_do; + + const ElementAcc* ptr_lse; + RowTensorStride stride_lse; + + const ElementAcc* ptr_sum_odo; + RowTensorStride stride_sum_odo; + + ElementAcc* ptr_dq_acc; + TensorStride stride_dq_acc; + + ElementAcc softmax_scale = 1.0f / sqrtf(TileShapeDQK{}); + }; + + using TMA_K = typename CollectiveMmaKQ::Params::TMA_A; + using TMA_V = typename CollectiveMmaVDO::Params::TMA_A; + using TMA_Q = typename CollectiveMmaKQ::Params::TMA_B; + using TMA_DO = typename CollectiveMmaVDO::Params::TMA_B; + + using TMA_DQ = decltype(make_tma_copy( + SM90_TMA_REDUCE_ADD{}, + make_tensor( + (const ElementAcc*)nullptr, + make_shape(1, 1, make_shape(make_shape(1, 1), 1)), + TensorStride{}), + SmemLayoutDQ{}(_, _, _0{}))); + + struct MainloopParams { + TMA_K tma_load_k; + TMA_V tma_load_v; + TMA_Q tma_load_q; + TMA_DO tma_load_do; + TMA_DQ tma_red_dq; + }; + + struct EpilogueArguments { + Element* ptr_dk; + TensorStride_GQA stride_dk; + Element* ptr_dv; + TensorStride_GQA stride_dv; + }; + + struct FnaArguments { + NADim q_shape; + NADim kv_shape; + NADim qkv_shape; + NADim window_size; + NADim stride; + NADim dilation; + }; + + struct FnaParams { + NADim qkv_shape; + NADim q_shape; + NADim kv_shape; + cute::tuple + na_params; // win, win_left, win_right, stride, stride_offset + bool is_fully_block_sparse; + bool has_q_padding; + NADim dilation; + bool requires_qkv_fixup; + bool is_dilated; + int num_dilation_groups; + }; + + struct Arguments { + ProblemShape problem_shape; + FnaArguments fna; + MainloopArguments mainloop; + EpilogueArguments epilogue; + KernelHardwareInfo hw_info; + }; + + struct Params { + ProblemShape problem_shape; + FnaParams fna; + MainloopArguments mainloop; + MainloopParams mainloop_params; + EpilogueArguments epilogue; + KernelHardwareInfo hw_info; + }; + + static bool can_implement(Arguments const& args) { + auto [Q, K, D, D_VO, HB] = args.problem_shape; + auto [H, B] = HB; + auto [H_R, H_K] = H; + if (Q <= 0 || K <= 0 || D <= 0 || D_VO <= 0 || H_R <= 0 || H_K <= 0 || + B <= 0) { + return false; + } + if (D % Alignment != 0 || D_VO % Alignment != 0) { + return false; + } + return evenly_divides(args.fna.q_shape, QTileShape{}) && + evenly_divides(args.fna.kv_shape, KVTileShape{}) && + evenly_divides(B, size(args.fna.dilation)) && // dilation groups are + // folded into batch + tuple_leq(args.fna.window_size, args.fna.qkv_shape) && + tuple_leq(args.fna.stride, args.fna.window_size); + } + + static Status initialize_workspace(Arguments const&, void*, cudaStream_t) { + return Status::kSuccess; + } + + static Params to_underlying_arguments(Arguments const& args, void*) { + auto [Q_, K_, D, D_VO, HB] = args.problem_shape; + int Q = Q_; + int K = K_; + + auto params_kq = CollectiveMmaKQ::to_underlying_arguments( + make_shape(K, Q, D, HB), + typename CollectiveMmaKQ::Arguments{ + args.mainloop.ptr_k, + args.mainloop.stride_k, + args.mainloop.ptr_q, + args.mainloop.stride_q, + }, + /*workspace=*/nullptr); + + auto params_vdo = CollectiveMmaVDO::to_underlying_arguments( + make_shape(K, Q, D_VO, HB), + typename CollectiveMmaVDO::Arguments{ + args.mainloop.ptr_v, + args.mainloop.stride_v, + args.mainloop.ptr_do, + args.mainloop.stride_do, + }, + /*workspace=*/nullptr); + + TMA_DQ tma_red_dq = make_tma_copy( + SM90_TMA_REDUCE_ADD{}, + make_tensor( + args.mainloop.ptr_dq_acc, + make_shape(Q_, D, HB), + args.mainloop.stride_dq_acc), + SmemLayoutDQ{}(_, _, _0{})); + + // FNA params + auto window_left = get_window_left(args.fna.window_size); + auto window_right = get_window_right(args.fna.window_size); + auto stride_offset = get_bwd_stride_offset(args.fna.stride); + auto na_params = make_tuple( + args.fna.window_size, + window_left, + window_right, + args.fna.stride, + stride_offset); + bool requires_qkv_fixup = + not evenly_divides(args.fna.qkv_shape, args.fna.dilation); + + auto is_fully_block_sparse = fully_block_sparse( + args.fna.qkv_shape, + args.fna.window_size, + args.fna.stride, + QTileShape{}, + KVTileShape{}); + + return Params{ + args.problem_shape, + FnaParams{ + args.fna.qkv_shape, + args.fna.q_shape, + args.fna.kv_shape, + na_params, + is_fully_block_sparse, + /* has_q_padding */ + not evenly_divides(args.fna.qkv_shape, QTileShape{}), + args.fna.dilation, + requires_qkv_fixup, + is_dilated(args.fna.dilation), + size(args.fna.dilation)}, + args.mainloop, + MainloopParams{ + params_kq.tma_load_a, + params_vdo.tma_load_a, + params_kq.tma_load_b, + params_vdo.tma_load_b, + tma_red_dq}, + args.epilogue, + args.hw_info}; + } + + template + static CUTLASS_DEVICE auto quantize(T const& input) { + constexpr int AlignmentS = 4; + auto output = make_tensor(shape(input)); + auto input_vec = recast>(input); + auto output_vec = recast>(output); + + cutlass::NumericArrayConverter epilogue_op; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(input_vec); i++) { + output_vec(i) = epilogue_op(input_vec(i)); + } + + return output; + } + + template < + class BlkCoord, + class BlkOffset, + class ProblemShape_, + class IterToTileMap> + CUTLASS_DEVICE void load( + BlkCoord const& blk_coord, + BlkOffset const& blk_offset, + ProblemShape_ const& problem_shape, + int iter_start, + int iter_end, + int iter_count, + IterToTileMap const& iter_to_tile_map, + MainloopArguments const& mainloop_args, + MainloopParams const& mainloop_params, + TensorStorage& shared_tensors, + PipelineLoadMmaQ& pipeline_load_mma_q, + typename PipelineLoadMmaQ::PipelineState& + pipeline_load_mma_q_producer_state, + PipelineLoadMmaDO& pipeline_load_mma_do, + typename PipelineLoadMmaDO::PipelineState& + pipeline_load_mma_do_producer_state, + PipelineLoadComputeLSE& pipeline_load_compute_lse, + typename PipelineLoadComputeLSE::PipelineState& + pipeline_load_compute_lse_producer_state, + PipelineLoadComputeSumOdO& pipeline_load_compute_sum_odo, + typename PipelineLoadComputeSumOdO::PipelineState& + pipeline_load_compute_sum_odo_producer_state) { + auto [Q, K, D, D_VO, HB] = problem_shape; + int iter_index = iter_start; + + using X = Underscore; + + uint16_t mcast_mask = 0; + + auto mK_in = + mainloop_params.tma_load_k.get_tma_tensor(make_shape(K, D, HB)); + auto mV_in = + mainloop_params.tma_load_v.get_tma_tensor(make_shape(K, D_VO, HB)); + auto mQ_in = + mainloop_params.tma_load_q.get_tma_tensor(make_shape(Q, D, HB)); + auto mDO_in = + mainloop_params.tma_load_do.get_tma_tensor(make_shape(Q, D_VO, HB)); + + auto mK = domain_offset(select<1, 2, 4>(blk_offset), mK_in); + auto mV = domain_offset(select<1, 3, 4>(blk_offset), mV_in); + auto mQ = domain_offset(select<0, 2, 4>(blk_offset), mQ_in); + auto mDO = domain_offset(select<0, 3, 4>(blk_offset), mDO_in); + + auto gK = + local_tile(mK, TileShapeKQ{}, make_coord(_, _, _), Step<_1, X, _1>{}); + auto gQ = + local_tile(mQ, TileShapeKQ{}, make_coord(_, _, _), Step{}); + auto gV = + local_tile(mV, TileShapeVDO{}, make_coord(_, _, _), Step<_1, X, _1>{}); + auto gDO = + local_tile(mDO, TileShapeVDO{}, make_coord(_, _, _), Step{}); + + ThrMMA cta_mma_kq = TiledMmaKQ{}.get_slice(_0{}); + ThrMMA cta_mma_vdo = TiledMmaVDO{}.get_slice(_0{}); + + auto tSTgK = cta_mma_kq.partition_A(gK); + auto tSTgQ = cta_mma_kq.partition_B(gQ); + auto tDPTgV = cta_mma_vdo.partition_A(gV); + auto tDPTgDO = cta_mma_vdo.partition_B(gDO); + + auto sQ = make_tensor( + make_smem_ptr(shared_tensors.smem_q.begin()), SmemLayoutQ{}); + auto sK = make_tensor( + make_smem_ptr(shared_tensors.smem_k.begin()), SmemLayoutK{}); + auto sV = make_tensor( + make_smem_ptr(shared_tensors.smem_v.begin()), SmemLayoutV{}); + auto sDO = make_tensor( + make_smem_ptr(shared_tensors.smem_do.begin()), SmemLayoutDO{}); + + auto [tKgK_mkl, tKsK] = tma_partition( + mainloop_params.tma_load_k, + _0{}, + make_layout(_1{}), + group_modes<0, 3>(sK), + group_modes<0, 3>(tSTgK)); + auto [tQgQ_mkl, tQsQ] = tma_partition( + mainloop_params.tma_load_q, + _0{}, + make_layout(_1{}), + group_modes<0, 3>(sQ), + group_modes<0, 3>(tSTgQ)); + auto [tVgV_mkl, tVsV] = tma_partition( + mainloop_params.tma_load_v, + _0{}, + make_layout(_1{}), + group_modes<0, 3>(sV), + group_modes<0, 3>(tDPTgV)); + auto [tDOgDO_mkl, tDOsDO] = tma_partition( + mainloop_params.tma_load_do, + _0{}, + make_layout(_1{}), + group_modes<0, 3>(sDO), + group_modes<0, 3>(tDPTgDO)); + + // set up lse and sum_odo + + auto + [blk_coord_q, blk_coord_k, blk_coord_d, blk_coord_dv, blk_coord_batch] = + blk_coord; + + pipeline_load_mma_q.producer_acquire(pipeline_load_mma_q_producer_state); + auto tma_barrier = pipeline_load_mma_q.producer_get_barrier( + pipeline_load_mma_q_producer_state); + + pipeline_load_mma_q.producer_expect_transaction( + pipeline_load_mma_q_producer_state, kTransactionsBytesLoadK); + + // load K + if (cute::elect_one_sync()) { + cute::copy( + mainloop_params.tma_load_k.with(*tma_barrier, mcast_mask), + tKgK_mkl(_, blk_coord_k, _0{}, blk_coord_batch), + tKsK(_, _0{})); + } + + // load Q + if (cute::elect_one_sync()) { + cute::copy( + mainloop_params.tma_load_q.with(*tma_barrier, mcast_mask), + tQgQ_mkl(_, iter_to_tile_map(iter_index), _0{}, blk_coord_batch), + tQsQ(_, pipeline_load_mma_q_producer_state.index())); + } + + ++pipeline_load_mma_q_producer_state; + + pipeline_load_compute_lse.producer_acquire( + pipeline_load_compute_lse_producer_state); + + // load LSE + // 32 threads loading 128 values of 32b each + // so 4*32b=128b + + int thread_idx = threadIdx.x % NumThreadsPerWarp; + int smem_idx = + TileShapeQ{} * pipeline_load_compute_lse_producer_state.index() + + thread_idx * 4; + int gmem_idx = TileShapeQ{} * iter_to_tile_map(iter_index) + thread_idx * 4; + auto mLSE = make_tensor( + mainloop_args.ptr_lse, make_shape(Q, HB), mainloop_args.stride_lse); + for (int i = 0; i < 4; i++) { + cutlass::arch::cp_async_zfill<4>( + shared_tensors.smem_lse.begin() + smem_idx + i, + &mLSE(gmem_idx + i, blk_coord_batch), + gmem_idx + i < Q); + } + + pipeline_load_compute_lse.producer_commit( + pipeline_load_compute_lse_producer_state, + cutlass::arch::cpasync_barrier_arrive); + ++pipeline_load_compute_lse_producer_state; + + pipeline_load_mma_do.producer_acquire(pipeline_load_mma_do_producer_state); + tma_barrier = pipeline_load_mma_do.producer_get_barrier( + pipeline_load_mma_do_producer_state); + + pipeline_load_mma_do.producer_expect_transaction( + pipeline_load_mma_do_producer_state, kTransactionsBytesLoadV); + + // load V + if (cute::elect_one_sync()) { + cute::copy( + mainloop_params.tma_load_v.with(*tma_barrier, mcast_mask), + tVgV_mkl(_, blk_coord_k, _0{}, blk_coord_batch), + tVsV(_, _0{})); + } + + // load dO + if (cute::elect_one_sync()) { + cute::copy( + mainloop_params.tma_load_do.with(*tma_barrier, mcast_mask), + tDOgDO_mkl(_, iter_to_tile_map(iter_index), _0{}, blk_coord_batch), + tDOsDO(_, pipeline_load_mma_do_producer_state.index())); + } + + ++pipeline_load_mma_do_producer_state; + + pipeline_load_compute_sum_odo.producer_acquire( + pipeline_load_compute_sum_odo_producer_state); + + // load sum_OdO + smem_idx = + TileShapeQ{} * pipeline_load_compute_sum_odo_producer_state.index() + + thread_idx * 4; + gmem_idx = TileShapeQ{} * iter_to_tile_map(iter_index) + thread_idx * 4; + auto mSumOdO = make_tensor( + mainloop_args.ptr_sum_odo, + make_shape(Q, HB), + mainloop_args.stride_sum_odo); + for (int i = 0; i < 4; i++) { + cutlass::arch::cp_async_zfill<4>( + shared_tensors.smem_sum_odo.begin() + smem_idx + i, + &mSumOdO(gmem_idx + i, blk_coord_batch), + gmem_idx + i < Q); + } + + pipeline_load_compute_sum_odo.producer_commit( + pipeline_load_compute_sum_odo_producer_state, + cutlass::arch::cpasync_barrier_arrive); + ++pipeline_load_compute_sum_odo_producer_state; + + iter_count -= 1; + iter_index += 1; + + while (iter_count > 0) { + if (iter_index == iter_end) { + iter_index = iter_start; + get<0, 0>(blk_coord_batch) += 1; + } + + pipeline_load_mma_q.producer_acquire(pipeline_load_mma_q_producer_state); + tma_barrier = pipeline_load_mma_q.producer_get_barrier( + pipeline_load_mma_q_producer_state); + + // load Q + if (cute::elect_one_sync()) { + cute::copy( + mainloop_params.tma_load_q.with(*tma_barrier, mcast_mask), + tQgQ_mkl(_, iter_to_tile_map(iter_index), _0{}, blk_coord_batch), + tQsQ(_, pipeline_load_mma_q_producer_state.index())); + } + + ++pipeline_load_mma_q_producer_state; + + pipeline_load_compute_lse.producer_acquire( + pipeline_load_compute_lse_producer_state); + + // load LSE + smem_idx = + TileShapeQ{} * pipeline_load_compute_lse_producer_state.index() + + thread_idx * 4; + gmem_idx = TileShapeQ{} * iter_to_tile_map(iter_index) + thread_idx * 4; + for (int i = 0; i < 4; i++) { + cutlass::arch::cp_async_zfill<4>( + shared_tensors.smem_lse.begin() + smem_idx + i, + &mLSE(gmem_idx + i, blk_coord_batch), + gmem_idx + i < Q); + } + + pipeline_load_compute_lse.producer_commit( + pipeline_load_compute_lse_producer_state, + cutlass::arch::cpasync_barrier_arrive); + ++pipeline_load_compute_lse_producer_state; + + pipeline_load_mma_do.producer_acquire( + pipeline_load_mma_do_producer_state); + tma_barrier = pipeline_load_mma_do.producer_get_barrier( + pipeline_load_mma_do_producer_state); + + // load dO + if (cute::elect_one_sync()) { + cute::copy( + mainloop_params.tma_load_do.with(*tma_barrier, mcast_mask), + tDOgDO_mkl(_, iter_to_tile_map(iter_index), _0{}, blk_coord_batch), + tDOsDO(_, pipeline_load_mma_do_producer_state.index())); + } + + ++pipeline_load_mma_do_producer_state; + + pipeline_load_compute_sum_odo.producer_acquire( + pipeline_load_compute_sum_odo_producer_state); + + // load sum_OdO + smem_idx = + TileShapeQ{} * pipeline_load_compute_sum_odo_producer_state.index() + + thread_idx * 4; + gmem_idx = TileShapeQ{} * iter_to_tile_map(iter_index) + thread_idx * 4; + for (int i = 0; i < 4; i++) { + cutlass::arch::cp_async_zfill<4>( + shared_tensors.smem_sum_odo.begin() + smem_idx + i, + &mSumOdO(gmem_idx + i, blk_coord_batch), + gmem_idx + i < Q); + } + + pipeline_load_compute_sum_odo.producer_commit( + pipeline_load_compute_sum_odo_producer_state, + cutlass::arch::cpasync_barrier_arrive); + ++pipeline_load_compute_sum_odo_producer_state; + + iter_count -= 1; + iter_index += 1; + } + } + + template + CUTLASS_DEVICE void mma( + BlkCoord const& blk_coord, + ProblemShape_ const& problem_shape, + int iter_start, + int iter_end, + int iter_count, + MainloopArguments const& mainloop_args, + TensorStorage& shared_tensors, + PipelineLoadMmaQ& pipeline_load_mma_q, + typename PipelineLoadMmaQ::PipelineState& + pipeline_load_mma_q_consumer_state, + PipelineLoadMmaDO& pipeline_load_mma_do, + typename PipelineLoadMmaDO::PipelineState& + pipeline_load_mma_do_consumer_state, + PipelineMmaComputeS& pipeline_mma_compute_s, + typename PipelineMmaComputeS::PipelineState& + pipeline_mma_compute_s_producer_state, + PipelineMmaComputeDP& pipeline_mma_compute_dp, + typename PipelineMmaComputeDP::PipelineState& + pipeline_mma_compute_dp_producer_state, + PipelineMmaReduceDQ& pipeline_mma_reduce_dq, + typename PipelineMmaReduceDQ::PipelineState& + pipeline_mma_reduce_dq_producer_state, + PipelineComputeMmaP& pipeline_compute_mma_p, + typename PipelineComputeMmaP::PipelineState& + pipeline_compute_mma_p_consumer_state, + PipelineComputeMmaDS& pipeline_compute_mma_ds, + typename PipelineComputeMmaDS::PipelineState& + pipeline_compute_mma_ds_consumer_state, + PipelineMmaComputeDKDV& pipeline_mma_compute_dkdv, + typename PipelineMmaComputeDKDV::PipelineState& + pipeline_mma_compute_dkdv_producer_state) { + auto [Q, K, D, D_VO, HB] = problem_shape; + + auto sQ = make_tensor( + make_smem_ptr(shared_tensors.smem_q.begin()), SmemLayoutQ{}); + auto sK = make_tensor( + make_smem_ptr(shared_tensors.smem_k.begin()), SmemLayoutK{}); + auto sV = make_tensor( + make_smem_ptr(shared_tensors.smem_v.begin()), SmemLayoutV{}); + auto sDO = make_tensor( + make_smem_ptr(shared_tensors.smem_do.begin()), SmemLayoutDO{}); + + auto sQT = make_tensor( + make_smem_ptr(shared_tensors.smem_q_t.begin()), SmemLayoutQT{}); + auto sKT = make_tensor( + make_smem_ptr(shared_tensors.smem_k_t.begin()), SmemLayoutKT{}); + auto sDS = make_tensor( + make_smem_ptr(shared_tensors.smem_ds.begin()), SmemLayoutDS{}); + auto sDST = make_tensor( + make_smem_ptr(shared_tensors.smem_ds_t.begin()), SmemLayoutDST{}); + auto sP = make_tensor( + make_smem_ptr((Element*)nullptr), + typename CollectiveMmaPDO::SmemLayoutA{}); + auto sDOT = make_tensor( + make_smem_ptr(shared_tensors.smem_do_t.begin()), SmemLayoutDOT{}); + + Tensor tSTrK = TiledMmaKQ::make_fragment_A(sK); + Tensor tSTrQ = TiledMmaKQ::make_fragment_B(sQ); + + Tensor tDPTrV = TiledMmaVDO::make_fragment_A(sV); + Tensor tDPTrDO = TiledMmaVDO::make_fragment_B(sDO); + + Tensor tDQrDS = TiledMmaDSK::make_fragment_A(sDS); + Tensor tDQrKT = TiledMmaDSK::make_fragment_B(sKT); + + Tensor tDKrDST = TiledMmaDSQ::make_fragment_A(sDST); + Tensor tDKrQT = TiledMmaDSQ::make_fragment_B(sQT); + + Tensor tDVrP = TiledMmaPDO::make_fragment_A(sP)(_, _, _, _0{}); + tDVrP.data() = TmemAllocation::kP; + Tensor tDVrDOT = TiledMmaPDO::make_fragment_B(sDOT); + + TiledMmaKQ tiled_mma_kq; + TiledMmaVDO tiled_mma_vdo; + TiledMmaDSK tiled_mma_dsk; + TiledMmaDSQ tiled_mma_dsq; + TiledMmaPDO tiled_mma_pdo; + + tiled_mma_dsq.accumulate_ = UMMA::ScaleOut::Zero; + tiled_mma_pdo.accumulate_ = UMMA::ScaleOut::Zero; + + Tensor tSTtST = + partition_fragment_C(tiled_mma_kq, select<0, 1>(TileShapeKQ{})); + tSTtST.data() = TmemAllocation::kS; + + Tensor tDPTtDPT = + partition_fragment_C(tiled_mma_vdo, select<0, 1>(TileShapeVDO{})); + tDPTtDPT.data() = TmemAllocation::kDP; + + Tensor tDQtDQ = + partition_fragment_C(tiled_mma_dsk, select<0, 1>(TileShapeDSK{})); + tDQtDQ.data() = TmemAllocation::kDQ; + + Tensor tDKtDK = + partition_fragment_C(tiled_mma_dsq, select<0, 1>(TileShapeDSQ{})); + tDKtDK.data() = TmemAllocation::kDK; + + Tensor tDVtDV = + partition_fragment_C(tiled_mma_pdo, select<0, 1>(TileShapePDO{})); + tDVtDV.data() = TmemAllocation::kDV; + + auto pipeline_load_mma_q_release_state = pipeline_load_mma_q_consumer_state; + + pipeline_load_mma_q.consumer_wait(pipeline_load_mma_q_consumer_state); + pipeline_mma_compute_s.producer_acquire( + pipeline_mma_compute_s_producer_state); + + // S = Q*K + tiled_mma_kq.accumulate_ = UMMA::ScaleOut::Zero; + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tSTrQ); ++k_block) { + cute::gemm( + tiled_mma_kq, + tSTrK(_, _, k_block, _0{}), + tSTrQ(_, _, k_block, pipeline_load_mma_q_consumer_state.index()), + tSTtST); + tiled_mma_kq.accumulate_ = UMMA::ScaleOut::One; + } + + ++pipeline_load_mma_q_consumer_state; + + pipeline_mma_compute_s.producer_commit( + pipeline_mma_compute_s_producer_state); + ++pipeline_mma_compute_s_producer_state; + + pipeline_load_mma_do.consumer_wait(pipeline_load_mma_do_consumer_state); + + pipeline_mma_compute_dp.producer_acquire( + pipeline_mma_compute_dp_producer_state); + pipeline_mma_reduce_dq.producer_acquire( + pipeline_mma_reduce_dq_producer_state); + + // dP = dO*V + tiled_mma_vdo.accumulate_ = UMMA::ScaleOut::Zero; + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tDPTrV); ++k_block) { + cute::gemm( + tiled_mma_vdo, + tDPTrV(_, _, k_block, _0{}), + tDPTrDO(_, _, k_block, pipeline_load_mma_do_consumer_state.index()), + tDPTtDPT); + tiled_mma_vdo.accumulate_ = UMMA::ScaleOut::One; + } + + pipeline_mma_compute_dp.producer_commit( + pipeline_mma_compute_dp_producer_state); + ++pipeline_mma_compute_dp_producer_state; + + pipeline_compute_mma_p.consumer_wait(pipeline_compute_mma_p_consumer_state); + + // dV = P*dO + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tDVrP); ++k_block) { + cute::gemm( + tiled_mma_pdo, + tDVrP(_, _, k_block), + tDVrDOT(_, _, k_block, pipeline_load_mma_do_consumer_state.index()), + tDVtDV); + tiled_mma_pdo.accumulate_ = UMMA::ScaleOut::One; + } + + pipeline_compute_mma_p.consumer_release( + pipeline_compute_mma_p_consumer_state); + ++pipeline_compute_mma_p_consumer_state; + + pipeline_load_mma_do.consumer_release(pipeline_load_mma_do_consumer_state); + ++pipeline_load_mma_do_consumer_state; + + iter_count -= 1; + + // in tmem, S & P overlap + // and dP and dQ overlap + // so we need to acquire dQ and dP at the same time + while (iter_count > 0) { + pipeline_load_mma_q.consumer_wait(pipeline_load_mma_q_consumer_state); + pipeline_mma_compute_s.producer_acquire( + pipeline_mma_compute_s_producer_state); + + // S = Q*K + tiled_mma_kq.accumulate_ = UMMA::ScaleOut::Zero; + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tSTrQ); ++k_block) { + cute::gemm( + tiled_mma_kq, + tSTrK(_, _, k_block, _0{}), + tSTrQ(_, _, k_block, pipeline_load_mma_q_consumer_state.index()), + tSTtST); + tiled_mma_kq.accumulate_ = UMMA::ScaleOut::One; + } + + ++pipeline_load_mma_q_consumer_state; + + pipeline_mma_compute_s.producer_commit( + pipeline_mma_compute_s_producer_state); + ++pipeline_mma_compute_s_producer_state; + + pipeline_compute_mma_ds.consumer_wait( + pipeline_compute_mma_ds_consumer_state); + + // we need to acquire dP here, because tmem dQ == tmem dP + pipeline_mma_compute_dp.producer_acquire( + pipeline_mma_compute_dp_producer_state); + + // dQ = dS*K + tiled_mma_dsk.accumulate_ = UMMA::ScaleOut::Zero; + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tDQrDS); ++k_block) { + cute::gemm( + tiled_mma_dsk, + tDQrDS( + _, _, k_block, pipeline_compute_mma_ds_consumer_state.index()), + tDQrKT(_, _, k_block, _0{}), + tDQtDQ); + tiled_mma_dsk.accumulate_ = UMMA::ScaleOut::One; + } + + pipeline_mma_reduce_dq.producer_commit( + pipeline_mma_reduce_dq_producer_state); + ++pipeline_mma_reduce_dq_producer_state; + + // dK = dS*Q + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tDKrDST); ++k_block) { + cute::gemm( + tiled_mma_dsq, + tDKrDST( + _, _, k_block, pipeline_compute_mma_ds_consumer_state.index()), + tDKrQT(_, _, k_block, pipeline_load_mma_q_release_state.index()), + tDKtDK); + tiled_mma_dsq.accumulate_ = UMMA::ScaleOut::One; + } + + pipeline_load_mma_q.consumer_release(pipeline_load_mma_q_release_state); + ++pipeline_load_mma_q_release_state; + + pipeline_compute_mma_ds.consumer_release( + pipeline_compute_mma_ds_consumer_state); + ++pipeline_compute_mma_ds_consumer_state; + + // we grab dq here, because in tmem dq == dp + pipeline_mma_reduce_dq.producer_acquire( + pipeline_mma_reduce_dq_producer_state); + + pipeline_load_mma_do.consumer_wait(pipeline_load_mma_do_consumer_state); + + // dP = dO*V + tiled_mma_vdo.accumulate_ = UMMA::ScaleOut::Zero; + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tDPTrV); ++k_block) { + cute::gemm( + tiled_mma_vdo, + tDPTrV(_, _, k_block, _0{}), + tDPTrDO(_, _, k_block, pipeline_load_mma_do_consumer_state.index()), + tDPTtDPT); + tiled_mma_vdo.accumulate_ = UMMA::ScaleOut::One; + } + + pipeline_mma_compute_dp.producer_commit( + pipeline_mma_compute_dp_producer_state); + ++pipeline_mma_compute_dp_producer_state; + + pipeline_compute_mma_p.consumer_wait( + pipeline_compute_mma_p_consumer_state); + + // dV = P*dO + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tDVrP); ++k_block) { + cute::gemm( + tiled_mma_pdo, + tDVrP(_, _, k_block), + tDVrDOT(_, _, k_block, pipeline_load_mma_do_consumer_state.index()), + tDVtDV); + tiled_mma_pdo.accumulate_ = UMMA::ScaleOut::One; + } + + pipeline_compute_mma_p.consumer_release( + pipeline_compute_mma_p_consumer_state); + ++pipeline_compute_mma_p_consumer_state; + + pipeline_load_mma_do.consumer_release( + pipeline_load_mma_do_consumer_state); + ++pipeline_load_mma_do_consumer_state; + + iter_count -= 1; + } + + // signal to the epilogue that dV is ready + pipeline_mma_compute_dkdv.producer_acquire( + pipeline_mma_compute_dkdv_producer_state); + pipeline_mma_compute_dkdv.producer_commit( + pipeline_mma_compute_dkdv_producer_state); + ++pipeline_mma_compute_dkdv_producer_state; + + pipeline_mma_compute_dkdv.producer_acquire( + pipeline_mma_compute_dkdv_producer_state); + + pipeline_compute_mma_ds.consumer_wait( + pipeline_compute_mma_ds_consumer_state); + + // dK = dS*Q + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tDKrDST); ++k_block) { + cute::gemm( + tiled_mma_dsq, + tDKrDST( + _, _, k_block, pipeline_compute_mma_ds_consumer_state.index()), + tDKrQT(_, _, k_block, pipeline_load_mma_q_release_state.index()), + tDKtDK); + tiled_mma_dsq.accumulate_ = UMMA::ScaleOut::One; + } + + // signal to epilgue that dK is ready + pipeline_mma_compute_dkdv.producer_commit( + pipeline_mma_compute_dkdv_producer_state); + ++pipeline_mma_compute_dkdv_producer_state; + + // we've already acquired mma_reduce_dq in the loop + + // dQ = dS*K + tiled_mma_dsk.accumulate_ = UMMA::ScaleOut::Zero; + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tDQrDS); ++k_block) { + cute::gemm( + tiled_mma_dsk, + tDQrDS(_, _, k_block, pipeline_compute_mma_ds_consumer_state.index()), + tDQrKT(_, _, k_block, _0{}), + tDQtDQ); + tiled_mma_dsk.accumulate_ = UMMA::ScaleOut::One; + } + + pipeline_mma_reduce_dq.producer_commit( + pipeline_mma_reduce_dq_producer_state); + ++pipeline_mma_reduce_dq_producer_state; + + pipeline_load_mma_q.consumer_release(pipeline_load_mma_q_release_state); + ++pipeline_load_mma_q_release_state; + + pipeline_compute_mma_ds.consumer_release( + pipeline_compute_mma_ds_consumer_state); + ++pipeline_compute_mma_ds_consumer_state; + } + + template + CUTLASS_DEVICE void store( + TensorG gmem, + TensorR const& regs, + TensorC const& coord, + TensorShape const& tensor_shape) { + Tensor preds = cute::lazy::transform( + coord, [&](auto const& c) { return elem_less(c, tensor_shape); }); + + auto copy_op = make_cotiled_copy( + Copy_Atom, Element>{}, + make_layout( + make_shape(_1{}, Int{})), + regs.layout()); + auto thr_copy = copy_op.get_slice(_0{}); + + Tensor quantized_regs = quantize(regs); + Tensor tCr = thr_copy.partition_S(quantized_regs); + Tensor tCg = thr_copy.partition_D(gmem); + Tensor tPc = thr_copy.partition_D(preds); + + copy_if(copy_op, tPc, tCr, tCg); + } + + template + CUTLASS_DEVICE void epilogue_clear( + BlkCoord const& blk_coord, + BlkOffset const& blk_offset, + ProblemShape_ const& problem_shape, + MainloopArguments const& mainloop_args, + EpilogueArguments const& epilogue_args) { + auto [Q, K, D, D_VO, HB] = problem_shape; + auto + [blk_coord_q, blk_coord_k, blk_coord_d, blk_coord_dv, blk_coord_batch] = + blk_coord; + + auto mDK_in = make_tensor( + make_gmem_ptr(epilogue_args.ptr_dk), + make_shape(K, TileShapeDQK{}, HB), + epilogue_args.stride_dk); + auto mDK = domain_offset(select<1, 2, 4>(blk_offset), mDK_in); + auto gDK = + local_tile(mDK, TileShapeDSQ{}, make_coord(_, _, _), Step<_1, _1, X>{})( + _, _, blk_coord_k, _0{}, blk_coord_batch); + + Tensor cDK = domain_offset( + make_coord(get<1>(blk_coord) * TileShapeK{}, _0{}), + make_identity_tensor(take<0, 2>(TileShapeDSQ{}))); + + auto mDV_in = make_tensor( + make_gmem_ptr(epilogue_args.ptr_dv), + make_shape(K, TileShapeDVO{}, HB), + epilogue_args.stride_dv); + auto mDV = domain_offset(select<1, 3, 4>(blk_offset), mDV_in); + auto gDV = + local_tile(mDV, TileShapePDO{}, make_coord(_, _, _), Step<_1, _1, X>{})( + _, _, blk_coord_k, _0{}, blk_coord_batch); + + Tensor cDV = domain_offset( + make_coord(blk_coord_k * TileShapeK{}, _0{}), + make_identity_tensor(take<0, 2>(TileShapePDO{}))); + + for (int i = threadIdx.x; i < size(gDK); i += blockDim.x) { + if (elem_less(cDK(i), select<1, 2>(problem_shape))) { + gDK(i) = Element(0); + } + } + for (int i = threadIdx.x; i < size(gDV); i += blockDim.x) { + if (elem_less(cDV(i), select<1, 3>(problem_shape))) { + gDV(i) = Element(0); + } + } + } + + template + CUTLASS_DEVICE void epilogue( + BlkCoord const& blk_coord, + BlkOffset const& blk_offset, + ProblemShape_ const& problem_shape, + MainloopArguments const& mainloop_args, + EpilogueArguments const& epilogue_args, + PipelineMmaComputeDKDV& pipeline_mma_compute_dkdv, + typename PipelineMmaComputeDKDV::PipelineState& + pipeline_mma_compute_dkdv_consumer_state) { + auto [Q, K, D, D_VO, HB] = problem_shape; + auto + [blk_coord_q, blk_coord_k, blk_coord_d, blk_coord_dv, blk_coord_batch] = + blk_coord; + + auto load_op = SM100_TMEM_LOAD_32dp32b16x{}; + + auto tDKtDK = + partition_fragment_C(TiledMmaDSQ{}, select<0, 1>(TileShapeDSQ{}))( + make_coord(_, _), _0{}, _0{}); + tDKtDK.data() = TmemAllocation::kDK; + + auto mDK_in = make_tensor( + make_gmem_ptr(epilogue_args.ptr_dk), + make_shape(K, TileShapeDQK{}, HB), + epilogue_args.stride_dk); + auto mDK = domain_offset(select<1, 2, 4>(blk_offset), mDK_in); + auto gDK = + local_tile(mDK, TileShapeDSQ{}, make_coord(_, _, _), Step<_1, _1, X>{})( + _, _, blk_coord_k, _0{}, blk_coord_batch); + + Tensor cDK = domain_offset( + make_coord(get<1>(blk_coord) * TileShapeK{}, _0{}), + make_identity_tensor(take<0, 2>(TileShapeDSQ{}))); + + constexpr int kNumWarpgroups = kNumComputeWarps / 4; + int dp_idx = threadIdx.x % 128; + int wg_idx = (threadIdx.x % (kNumComputeWarps * NumThreadsPerWarp)) / 128; + + auto split_wg = [&](auto const& t) { + if constexpr (decltype(rank(t))::value == 3) { + auto p = t.compose(make_layout(make_shape( + size<0>(t), + size<1>(t), + make_shape( + Int{}, size<2>(t) / Int{})))); + return p(_, _, make_coord(wg_idx, _)); + } else { + auto p = t.compose(make_layout(make_shape( + size<0>(t), + size<1>(t), + size<2>(t), + make_shape( + Int{}, size<3>(t) / Int{})))); + return p(_, _, _, make_coord(wg_idx, _)); + } + }; + + auto tiled_t2r_dk = make_tmem_copy(load_op, tDKtDK); + auto thread_t2r_dk = tiled_t2r_dk.get_slice(dp_idx); + + Tensor tTR_cDK = split_wg(thread_t2r_dk.partition_D(cDK)); + Tensor tTR_gDK = split_wg(thread_t2r_dk.partition_D(gDK)); + Tensor tTR_rDK = make_tensor(shape(tTR_cDK)); + Tensor tTR_tDK = split_wg(thread_t2r_dk.partition_S(tDKtDK)); + + auto tDVtDV = + partition_fragment_C(TiledMmaDSQ{}, select<0, 1>(TileShapeDSQ{}))( + make_coord(_, _), _0{}, _0{}); + tDVtDV.data() = TmemAllocation::kDV; + + auto mDV_in = make_tensor( + make_gmem_ptr(epilogue_args.ptr_dv), + make_shape(K, TileShapeDVO{}, HB), + epilogue_args.stride_dv); + auto mDV = domain_offset(select<1, 3, 4>(blk_offset), mDV_in); + auto gDV = + local_tile(mDV, TileShapePDO{}, make_coord(_, _, _), Step<_1, _1, X>{})( + _, _, blk_coord_k, _0{}, blk_coord_batch); + + Tensor cDV = domain_offset( + make_coord(blk_coord_k * TileShapeK{}, _0{}), + make_identity_tensor(take<0, 2>(TileShapePDO{}))); + + auto tiled_t2r_dv = make_tmem_copy(load_op, tDVtDV); + auto thread_t2r_dv = tiled_t2r_dv.get_slice(dp_idx); + + Tensor tTR_cDV = split_wg(thread_t2r_dv.partition_D(cDV)); + Tensor tTR_gDV = split_wg(thread_t2r_dv.partition_D(gDV)); + Tensor tTR_rDV = make_tensor(shape(tTR_cDV)); + Tensor tTR_tDV = split_wg(thread_t2r_dv.partition_S(tDVtDV)); + + pipeline_mma_compute_dkdv.consumer_wait( + pipeline_mma_compute_dkdv_consumer_state); + + // load tDVtDV + cute::copy(tiled_t2r_dv, tTR_tDV, tTR_rDV); + + // store tDVgDV + store(tTR_gDV, tTR_rDV, tTR_cDV, select<1, 3>(problem_shape)); + + cutlass::arch::fence_view_async_tmem_load(); + pipeline_mma_compute_dkdv.consumer_release( + pipeline_mma_compute_dkdv_consumer_state); + ++pipeline_mma_compute_dkdv_consumer_state; + + pipeline_mma_compute_dkdv.consumer_wait( + pipeline_mma_compute_dkdv_consumer_state); + + // load tDKtDK + cute::copy(tiled_t2r_dk, tTR_tDK, tTR_rDK); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tTR_rDK); i++) { + tTR_rDK(i) = mainloop_args.softmax_scale * tTR_rDK(i); + } + + // store tDKgDK + store(tTR_gDK, tTR_rDK, tTR_cDK, select<1, 2>(problem_shape)); + + cutlass::arch::fence_view_async_tmem_load(); + pipeline_mma_compute_dkdv.consumer_release( + pipeline_mma_compute_dkdv_consumer_state); + ++pipeline_mma_compute_dkdv_consumer_state; + } + + template < + class BlkCoord, + class BlkOffset, + class ProblemShape_, + class QKVShape, + class NAParams> + CUTLASS_DEVICE void compute( + BlkCoord const& blk_coord, + BlkOffset const& blk_offset, + ProblemShape_ const& problem_shape, + int iter_start, + int iter_end, + int iter_count, + QKVShape const& kv_shape, + QKVShape const& qkv_shape, + NAParams const& na_params, + QKVShape const& q_offset, + QKVShape const& q_tiles, + bool is_fully_block_sparse, + bool has_q_padding, + MainloopArguments const& mainloop_args, + EpilogueArguments const& epilogue_args, + TensorStorage& shared_tensors, + PipelineLoadComputeLSE& pipeline_load_compute_lse, + typename PipelineLoadComputeLSE::PipelineState& + pipeline_load_compute_lse_consumer_state, + PipelineLoadComputeSumOdO& pipeline_load_compute_sum_odo, + typename PipelineLoadComputeSumOdO::PipelineState& + pipeline_load_compute_sum_odo_consumer_state, + PipelineMmaComputeS& pipeline_mma_compute_s, + typename PipelineMmaComputeS::PipelineState& + pipeline_mma_compute_s_consumer_state, + PipelineMmaComputeDP& pipeline_mma_compute_dp, + typename PipelineMmaComputeDP::PipelineState& + pipeline_mma_compute_dp_consumer_state, + PipelineComputeMmaP& pipeline_compute_mma_p, + typename PipelineComputeMmaP::PipelineState& + pipeline_compute_mma_p_producer_state, + PipelineComputeMmaDS& pipeline_compute_mma_ds, + typename PipelineComputeMmaDS::PipelineState& + pipeline_compute_mma_ds_producer_state, + PipelineMmaComputeDKDV& pipeline_mma_compute_dkdv, + typename PipelineMmaComputeDKDV::PipelineState& + pipeline_mma_compute_dkdv_consumer_state) { + auto [Q, K, D, D_VO, HB] = problem_shape; + int iter_index = iter_start; + + // in tmem, S & P overlap + // and dP and dQ overlap + + // there are two compute wg's that cooperatively compute softmax + // they are striped by this tmem atom, i.e. wg0 has 16 elems, then wg1 etc + + auto load_op = SM100_TMEM_LOAD_32dp32b16x{}; + auto store_op = []() { + if constexpr (sizeof(Element) == 1) { + return SM100_TMEM_STORE_32dp32b4x{}; + } else { + return SM100_TMEM_STORE_32dp32b8x{}; + } + }(); + + Tensor tSTtST = + partition_fragment_C(TiledMmaKQ{}, select<0, 1>(TileShapeKQ{}))( + make_coord(_, _), _0{}, _0{}); + tSTtST.data() = TmemAllocation::kS; + + Tensor tDPTtDPT = + partition_fragment_C(TiledMmaVDO{}, select<0, 1>(TileShapeVDO{}))( + make_coord(_, _), _0{}, _0{}); + tDPTtDPT.data() = TmemAllocation::kDP; + + Tensor cST = make_identity_tensor(take<0, 2>(TileShapeKQ{})); + Tensor cDPT = make_identity_tensor(take<0, 2>(TileShapeVDO{})); + + constexpr int kNumWarpgroups = kNumComputeWarps / 4; + int dp_idx = threadIdx.x % 128; + int wg_idx = (threadIdx.x % (kNumComputeWarps * NumThreadsPerWarp)) / 128; + auto tiled_t2r = make_tmem_copy(load_op, tSTtST); + auto thread_t2r = tiled_t2r.get_slice(dp_idx); + + auto split_wg = [&](auto const& t) { + if constexpr (decltype(size<1>(t))::value > 1) { + if constexpr (decltype(rank(t))::value == 3) { + auto p = t.compose(make_layout(make_shape( + size<0>(t), + make_shape( + Int{}, size<1>(t) / Int{}), + size<2>(t)))); + return p(_, make_coord(wg_idx, _), _); + } else { + auto p = t.compose(make_layout(make_shape( + size<0>(t), + make_shape( + Int{}, size<1>(t) / Int{}), + size<2>(t), + size<3>(t)))); + return p(_, make_coord(wg_idx, _), _, _); + } + } else { + if constexpr (decltype(rank(t))::value == 3) { + auto p = t.compose(make_layout(make_shape( + size<0>(t), + size<1>(t), + make_shape( + Int{}, size<2>(t) / Int{})))); + return p(_, _, make_coord(wg_idx, _)); + } else { + auto p = t.compose(make_layout(make_shape( + size<0>(t), + size<1>(t), + size<2>(t), + make_shape( + Int{}, size<3>(t) / Int{})))); + return p(_, _, _, make_coord(wg_idx, _)); + } + } + }; + + Tensor tTR_cST_p = thread_t2r.partition_D(cST); + Tensor tTR_cST = split_wg(tTR_cST_p); + Tensor tTR_rST = make_tensor(shape(tTR_cST)); + // Tensor tTR_tST_p = thread_t2r.partition_S(tSTtST); + Tensor tTR_tST = split_wg(thread_t2r.partition_S(tSTtST)); + + Tensor tTR_cDPT_p = thread_t2r.partition_D(cDPT); + Tensor tTR_cDPT = split_wg(tTR_cDPT_p); + Tensor tTR_rDPT = make_tensor(shape(tTR_cDPT)); + Tensor tTR_tDPT = split_wg(thread_t2r.partition_S(tDPTtDPT)); + + Tensor sLSE = make_tensor( + make_smem_ptr(shared_tensors.smem_lse.begin()), SmemLayoutLSE{}); + Tensor sSumOdO = make_tensor( + make_smem_ptr(shared_tensors.smem_sum_odo.begin()), SmemLayoutSumOdO{}); + + auto sP = make_tensor( + make_smem_ptr((Element*)nullptr), + typename CollectiveMmaPDO::SmemLayoutA{}); + + auto tDVrP = TiledMmaPDO::make_fragment_A(sP)(_, _, _, _0{}); + auto tDVcST = TiledMmaPDO{}.get_slice(_0{}).partition_A(cST); + tDVrP.data() = TmemAllocation::kP; + + auto tiled_r2t = make_tmem_copy(store_op, tDVrP); + auto thread_r2t = tiled_r2t.get_slice(dp_idx); + + auto tRT_tP = split_wg(thread_r2t.partition_D(tDVrP)); + auto tRT_cST_p = thread_r2t.partition_S(tDVcST); + auto tRT_cST = split_wg(tRT_cST_p); + + bool is_residual_k = + get<1>(blk_coord) * TileShapeK{} + TileShapeK{} > get<1>(problem_shape); + + CUTLASS_PRAGMA_NO_UNROLL + while (iter_count > 0) { + // wait for S and P + pipeline_mma_compute_s.consumer_wait( + pipeline_mma_compute_s_consumer_state); + pipeline_compute_mma_p.producer_acquire( + pipeline_compute_mma_p_producer_state); + // wait for LSE + pipeline_load_compute_lse.consumer_wait( + pipeline_load_compute_lse_consumer_state); + + // compute P = softmax(S, LSE) + cute::copy(tiled_t2r, tTR_tST, tTR_rST); + + if (not is_fully_block_sparse) { + Mask{}.apply_mask( + tTR_rST, + /* index_q = */ + [&](int i) { + auto c_transpose = tTR_cST(i); + // NOTE: we only care about q_idx here! + return get<1>(c_transpose) + iter_index * TileShapeQ{}; + }, + /* index_qk = */ + [&](int i) { + auto c_transpose = tTR_cST(i); + return make_coord( + get<1>(c_transpose) + iter_index * TileShapeQ{}, + get<0>(c_transpose) + get<1>(blk_coord) * TileShapeK{}); + }, + MultiDimTileShape{}, + kv_shape, + qkv_shape, + na_params, + q_offset, + q_tiles); + } else if (has_q_padding) { + Mask{}.apply_padded_mask( + tTR_rST, + /* index_q = */ + [&](int i) { + auto c_transpose = tTR_cST(i); + // NOTE: we only care about q_idx here! + return get<1>(c_transpose) + iter_index * TileShapeQ{}; + }, + MultiDimTileShape{}, + qkv_shape, + na_params, + q_offset, + q_tiles); + } + + ElementAcc log2_e = static_cast(M_LOG2E); + float2 softmax_scale_log2_e; + softmax_scale_log2_e.x = mainloop_args.softmax_scale * log2_e; + softmax_scale_log2_e.y = mainloop_args.softmax_scale * log2_e; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tTR_rST); i += 2) { + float2 acc; + float2 lse; + float2 out; + acc.x = tTR_rST(i); + acc.y = tTR_rST(i + 1); + lse.x = sLSE( + get<1>(tTR_cST(i)), + pipeline_load_compute_lse_consumer_state.index()); + lse.y = sLSE( + get<1>(tTR_cST(i + 1)), + pipeline_load_compute_lse_consumer_state.index()); + cute::fma(out, softmax_scale_log2_e, acc, lse); + tTR_rST(i) = ::exp2f(out.x); + tTR_rST(i + 1) = ::exp2f(out.y); + } + + auto tRT_rST = quantize(tTR_rST); + auto tRT_rST_reshaped = make_tensor(tRT_rST.data(), shape(tRT_cST)); + + cutlass::arch::fence_view_async_tmem_load(); + cutlass::arch::NamedBarrier( + kNumComputeWarps * NumThreadsPerWarp, + cutlass::arch::ReservedNamedBarriers::TransformBarrier) + .arrive_and_wait(); + + cute::copy(tiled_r2t, tRT_rST_reshaped, tRT_tP); + + // notify for P + cutlass::arch::fence_view_async_tmem_store(); + pipeline_compute_mma_p.producer_commit( + pipeline_compute_mma_p_producer_state); + ++pipeline_compute_mma_p_producer_state; + // release S + pipeline_mma_compute_s.consumer_release( + pipeline_mma_compute_s_consumer_state); + ++pipeline_mma_compute_s_consumer_state; + // release LSE + pipeline_load_compute_lse.consumer_release( + pipeline_load_compute_lse_consumer_state); + ++pipeline_load_compute_lse_consumer_state; + + // wait for OdO + pipeline_load_compute_sum_odo.consumer_wait( + pipeline_load_compute_sum_odo_consumer_state); + // wait for dP + pipeline_mma_compute_dp.consumer_wait( + pipeline_mma_compute_dp_consumer_state); + + // wait for dS + // in principle, we could defer waiting for dS, and move in the freeing of + // dP however, that would force us to keep dS in registers longer + pipeline_compute_mma_ds.producer_acquire( + pipeline_compute_mma_ds_producer_state); + + // compute dS = dsoftmax(P, dP, sum_OdO) + cute::copy(tiled_t2r, tTR_tDPT, tTR_rDPT); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tTR_rDPT); i += 2) { + float2 st; + st.x = tTR_rST(i); + st.y = tTR_rST(i + 1); + float2 dpt; + dpt.x = tTR_rDPT(i); + dpt.y = tTR_rDPT(i + 1); + float2 odo; + odo.x = sSumOdO( + get<1>(tTR_cDPT(i)), + pipeline_load_compute_sum_odo_consumer_state.index()); + odo.y = sSumOdO( + get<1>(tTR_cDPT(i + 1)), + pipeline_load_compute_sum_odo_consumer_state.index()); + float2 dif; + // sum odo is negated during preprocess + cute::add(dif, dpt, odo); + float2 out; + cute::mul(out, dif, st); + tTR_rDPT(i) = out.x; + tTR_rDPT(i + 1) = out.y; + } + + auto tTR_rDST = quantize(tTR_rDPT); + + // release dP + cutlass::arch::fence_view_async_tmem_load(); + pipeline_mma_compute_dp.consumer_release( + pipeline_mma_compute_dp_consumer_state); + ++pipeline_mma_compute_dp_consumer_state; + + Tensor sDS = make_tensor( + make_smem_ptr((Element*)shared_tensors.smem_ds.begin()), + SmemLayoutDS{})( + _, _, _, pipeline_compute_mma_ds_producer_state.index()); + + auto thread_layout = make_ordered_layout( + make_shape(_128{}, _128{}), make_stride(_1{}, _0{})); + + auto sDS_pi = as_position_independent_swizzle_tensor(sDS); + auto sDS_pi_slice_p = sDS_pi.compose(thread_layout)(dp_idx, _).compose( + make_layout(shape(tTR_cDPT_p))); + auto sDS_pi_slice = split_wg(sDS_pi_slice_p); + + copy_aligned(tTR_rDST, sDS_pi_slice); + + // notify for dS + cutlass::arch::fence_view_async_shared(); + pipeline_compute_mma_ds.producer_commit( + pipeline_compute_mma_ds_producer_state); + ++pipeline_compute_mma_ds_producer_state; + // release OdO + pipeline_load_compute_sum_odo.consumer_release( + pipeline_load_compute_sum_odo_consumer_state); + ++pipeline_load_compute_sum_odo_consumer_state; + + iter_count -= 1; + iter_index += 1; + if (iter_index == iter_end) { + iter_index = iter_start; + } + } + + epilogue( + blk_coord, + blk_offset, + problem_shape, + mainloop_args, + epilogue_args, + pipeline_mma_compute_dkdv, + pipeline_mma_compute_dkdv_consumer_state); + } + + template + CUTLASS_DEVICE void reduce( + BlkCoord const& blk_coord, + ProblemShape_ const& problem_shape, + int iter_start, + int iter_end, + int iter_count, + IterToTileMap const& iter_to_tile_map, + MainloopArguments const& mainloop_args, + MainloopParams const& mainloop_params, + TensorStorage& shared_tensors, + PipelineMmaReduceDQ& pipeline_mma_reduce_dq, + typename PipelineMmaReduceDQ::PipelineState& + pipeline_mma_reduce_dq_consumer_state, + PipelineReduceTmaStore& pipeline_reduce_tma_store, + typename PipelineReduceTmaStore::PipelineState& + pipeline_reduce_tma_store_producer_state) { + using X = Underscore; + + auto [Q, K, D, D_VO, HB] = problem_shape; + int iter_index = iter_start; + + auto + [blk_coord_q, blk_coord_k, blk_coord_d, blk_coord_dv, blk_coord_batch] = + blk_coord; + + // must match TileShapeDQ + auto load_op = SM100_TMEM_LOAD_32dp32b32x{}; + + auto tDQtDQ = + partition_fragment_C(TiledMmaDSK{}, select<0, 1>(TileShapeDSK{}))( + make_coord(_, _), _0{}, _0{}); + tDQtDQ.data() = TmemAllocation::kDQ; + + Tensor mDQ = + mainloop_params.tma_red_dq.get_tma_tensor(make_shape(Q, D, HB)); + auto gDQ = + local_tile(mDQ, TileShapeKQ{}, make_coord(_, _, _), Step{})( + _, _, _, _0{}, _); + + Tensor cDQ = make_identity_tensor(take<0, 2>(TileShapeDSK{})); + + Tensor sDQ = make_tensor( + make_smem_ptr(shared_tensors.smem_dq.begin()), SmemLayoutDQ{}); + + int thread_idx = threadIdx.x % (kNumComputeWarps * NumThreadsPerWarp); + auto tiled_t2r = make_tmem_copy(load_op, tDQtDQ); + auto thread_t2r = tiled_t2r.get_slice(thread_idx); + + Tensor tTR_cDQ = thread_t2r.partition_D(cDQ); + Tensor tTR_sDQ = thread_t2r.partition_D(sDQ); + Tensor tTR_tDQ = thread_t2r.partition_S(tDQtDQ); + + auto block_tma = mainloop_params.tma_red_dq.get_slice(_0{}); + + Tensor tDQsDQ = block_tma.partition_S(sDQ); + Tensor tDQcDQ = block_tma.partition_S(cDQ); + Tensor tDQgDQ = block_tma.partition_D(gDQ); + + int lane_predicate = + (threadIdx.x % (kNumReduceWarps * NumThreadsPerWarp)) == 0; + + while (iter_count > 0) { + pipeline_mma_reduce_dq.consumer_wait( + pipeline_mma_reduce_dq_consumer_state); + + Tensor tTR_rDQ = make_tensor(shape(tTR_cDQ)); + + // load dQ from tmem to rmem + cute::copy(tiled_t2r, tTR_tDQ, tTR_rDQ); + + cutlass::arch::fence_view_async_tmem_load(); + pipeline_mma_reduce_dq.consumer_release( + pipeline_mma_reduce_dq_consumer_state); + ++pipeline_mma_reduce_dq_consumer_state; + + // we don't have enough smem to dump it all to smem, so we do it in stages + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<2>(tTR_cDQ); i++) { + if (lane_predicate) { + pipeline_reduce_tma_store.producer_acquire( + pipeline_reduce_tma_store_producer_state); + } + // wait in all threads for the acquire to complete + cutlass::arch::NamedBarrier( + kNumReduceWarps * NumThreadsPerWarp, + cutlass::arch::ReservedNamedBarriers::TransposeBarrier) + .arrive_and_wait(); + + cute::copy( + tTR_rDQ(_, _, i), + tTR_sDQ( + _, _, _0{}, pipeline_reduce_tma_store_producer_state.index())); + + // wait for the stores to all be visible to the TMA + cutlass::arch::fence_view_async_shared(); + cutlass::arch::NamedBarrier( + kNumReduceWarps * NumThreadsPerWarp, + cutlass::arch::ReservedNamedBarriers::TransposeBarrier) + .arrive_and_wait(); + if (lane_predicate) { + // launch tma store + copy( + mainloop_params.tma_red_dq, + tDQsDQ( + _, _, _0{}, pipeline_reduce_tma_store_producer_state.index()), + tDQgDQ(_, _, i, iter_to_tile_map(iter_index), blk_coord_batch)); + pipeline_reduce_tma_store.producer_commit( + pipeline_reduce_tma_store_producer_state); + } + + ++pipeline_reduce_tma_store_producer_state; + } + + iter_count -= 1; + iter_index += 1; + if (iter_index == iter_end) { + iter_index = iter_start; + get<0, 0>(blk_coord_batch) += 1; + } + } + } + + CUTLASS_DEVICE void operator()(Params const& params, char* smem) { +#if ( \ + !defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) && \ + !defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) && \ + !defined(CUTLASS_ARCH_MMA_SM103A_ENABLED) && \ + !defined(CUTLASS_ARCH_MMA_SM103F_ENABLED)) + printf( + "ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); +#else + int warp_idx = cutlass::canonical_warp_idx_sync(); + auto role = warp_idx_to_role(warp_idx); + uint32_t lane_predicate = cute::elect_one_sync(); + + if (role == WarpRole::Load && lane_predicate) { + prefetch_tma_descriptor( + params.mainloop_params.tma_load_q.get_tma_descriptor()); + prefetch_tma_descriptor( + params.mainloop_params.tma_load_k.get_tma_descriptor()); + prefetch_tma_descriptor( + params.mainloop_params.tma_load_v.get_tma_descriptor()); + prefetch_tma_descriptor( + params.mainloop_params.tma_load_do.get_tma_descriptor()); + } + + SharedStorage& shared_storage = *reinterpret_cast(smem); + + int initializing_warp = 0; + typename PipelineLoadMmaQ::Params pipeline_load_mma_q_params; + if (role == WarpRole::Load) { + pipeline_load_mma_q_params.role = + PipelineLoadMmaQ::ThreadCategory::Producer; + } + if (role == WarpRole::Mma) { + pipeline_load_mma_q_params.role = + PipelineLoadMmaQ::ThreadCategory::Consumer; + } + pipeline_load_mma_q_params.is_leader = + lane_predicate && (role == WarpRole::Load); + // Also loads K in the first iteration + pipeline_load_mma_q_params.transaction_bytes = kTransactionsBytesLoadQ; + pipeline_load_mma_q_params.initializing_warp = initializing_warp++; + PipelineLoadMmaQ pipeline_load_mma_q( + shared_storage.pipelines.load_mma_q, + pipeline_load_mma_q_params, + ClusterShape{}, + /*barrier init*/ cute::true_type{}, + /*mask calc*/ cute::false_type{}); + + typename PipelineLoadMmaDO::Params pipeline_load_mma_do_params; + if (role == WarpRole::Load) { + pipeline_load_mma_do_params.role = + PipelineLoadMmaDO::ThreadCategory::Producer; + } + if (role == WarpRole::Mma) { + pipeline_load_mma_do_params.role = + PipelineLoadMmaDO::ThreadCategory::Consumer; + } + pipeline_load_mma_do_params.is_leader = + lane_predicate && (role == WarpRole::Load); + // Also loads V in the first iteration + pipeline_load_mma_do_params.transaction_bytes = kTransactionsBytesLoadDO; + pipeline_load_mma_do_params.initializing_warp = initializing_warp++; + PipelineLoadMmaDO pipeline_load_mma_do( + shared_storage.pipelines.load_mma_do, + pipeline_load_mma_do_params, + ClusterShape{}, + /*barrier init*/ cute::true_type{}, + /*mask calc*/ cute::false_type{}); + + typename PipelineLoadComputeLSE::Params pipeline_load_compute_lse_params; + if (role == WarpRole::Load) { + pipeline_load_compute_lse_params.role = + PipelineLoadComputeLSE::ThreadCategory::Producer; + } + if (role == WarpRole::Compute) { + pipeline_load_compute_lse_params.role = + PipelineLoadComputeLSE::ThreadCategory::Consumer; + } + pipeline_load_compute_lse_params.producer_arv_count = NumThreadsPerWarp; + pipeline_load_compute_lse_params.consumer_arv_count = + kNumComputeWarps * NumThreadsPerWarp; + pipeline_load_compute_lse_params.initializing_warp = initializing_warp++; + PipelineLoadComputeLSE pipeline_load_compute_lse( + shared_storage.pipelines.load_compute_lse, + pipeline_load_compute_lse_params, + /*barrier init*/ cute::true_type{}); + + typename PipelineLoadComputeSumOdO::Params + pipeline_load_compute_sum_odo_params; + if (role == WarpRole::Load) { + pipeline_load_compute_sum_odo_params.role = + PipelineLoadComputeSumOdO::ThreadCategory::Producer; + } + if (role == WarpRole::Compute) { + pipeline_load_compute_sum_odo_params.role = + PipelineLoadComputeSumOdO::ThreadCategory::Consumer; + } + pipeline_load_compute_sum_odo_params.producer_arv_count = NumThreadsPerWarp; + pipeline_load_compute_sum_odo_params.consumer_arv_count = + kNumComputeWarps * NumThreadsPerWarp; + pipeline_load_compute_sum_odo_params.initializing_warp = + initializing_warp++; + PipelineLoadComputeSumOdO pipeline_load_compute_sum_odo( + shared_storage.pipelines.load_compute_sum_odo, + pipeline_load_compute_sum_odo_params, + /*barrier init*/ cute::true_type{}); + + typename PipelineMmaComputeS::Params pipeline_mma_compute_s_params; + if (role == WarpRole::Mma) { + pipeline_mma_compute_s_params.role = + PipelineMmaComputeS::ThreadCategory::Producer; + } + if (role == WarpRole::Compute) { + pipeline_mma_compute_s_params.role = + PipelineMmaComputeS::ThreadCategory::Consumer; + } + pipeline_mma_compute_s_params.consumer_arv_count = + kNumComputeWarps * cutlass::NumThreadsPerWarp; + pipeline_mma_compute_s_params.initializing_warp = initializing_warp++; + PipelineMmaComputeS pipeline_mma_compute_s( + shared_storage.pipelines.mma_compute_s, + pipeline_mma_compute_s_params, + ClusterShape{}, + /*barrier init*/ cute::true_type{}, + /*mask calc*/ cute::false_type{}); + + typename PipelineMmaComputeDP::Params pipeline_mma_compute_dp_params; + if (role == WarpRole::Mma) { + pipeline_mma_compute_dp_params.role = + PipelineMmaComputeDP::ThreadCategory::Producer; + } + if (role == WarpRole::Compute) { + pipeline_mma_compute_dp_params.role = + PipelineMmaComputeDP::ThreadCategory::Consumer; + } + pipeline_mma_compute_dp_params.consumer_arv_count = + kNumComputeWarps * cutlass::NumThreadsPerWarp; + pipeline_mma_compute_dp_params.initializing_warp = initializing_warp++; + PipelineMmaComputeDP pipeline_mma_compute_dp( + shared_storage.pipelines.mma_compute_dp, + pipeline_mma_compute_dp_params, + ClusterShape{}, + /*barrier init*/ cute::true_type{}, + /*mask calc*/ cute::false_type{}); + + typename PipelineMmaReduceDQ::Params pipeline_mma_reduce_dq_params; + if (role == WarpRole::Mma) { + pipeline_mma_reduce_dq_params.role = + PipelineMmaReduceDQ::ThreadCategory::Producer; + } + if (role == WarpRole::Reduce) { + pipeline_mma_reduce_dq_params.role = + PipelineMmaReduceDQ::ThreadCategory::Consumer; + } + pipeline_mma_reduce_dq_params.consumer_arv_count = + kNumReduceWarps * cutlass::NumThreadsPerWarp; + pipeline_mma_reduce_dq_params.initializing_warp = initializing_warp++; + PipelineMmaReduceDQ pipeline_mma_reduce_dq( + shared_storage.pipelines.mma_reduce_dq, + pipeline_mma_reduce_dq_params, + ClusterShape{}, + /*barrier init*/ cute::true_type{}, + /*mask calc*/ cute::false_type{}); + + typename PipelineComputeMmaP::Params pipeline_compute_mma_p_params; + if (role == WarpRole::Mma) { + pipeline_compute_mma_p_params.role = + PipelineComputeMmaP::ThreadCategory::Consumer; + } + if (role == WarpRole::Compute) { + pipeline_compute_mma_p_params.role = + PipelineComputeMmaP::ThreadCategory::Producer; + } + pipeline_compute_mma_p_params.producer_arv_count = + kNumComputeWarps * cutlass::NumThreadsPerWarp; + pipeline_compute_mma_p_params.consumer_arv_count = 1; + pipeline_compute_mma_p_params.initializing_warp = initializing_warp++; + PipelineComputeMmaP pipeline_compute_mma_p( + shared_storage.pipelines.compute_mma_p, + pipeline_compute_mma_p_params, + ClusterShape{}, + /*barrier init*/ cute::true_type{}, + /*mask calc*/ cute::false_type{}); + + typename PipelineComputeMmaDS::Params pipeline_compute_mma_ds_params; + if (role == WarpRole::Mma) { + pipeline_compute_mma_ds_params.role = + PipelineComputeMmaDS::ThreadCategory::Consumer; + } + if (role == WarpRole::Compute) { + pipeline_compute_mma_ds_params.role = + PipelineComputeMmaDS::ThreadCategory::Producer; + } + pipeline_compute_mma_ds_params.producer_arv_count = + kNumComputeWarps * cutlass::NumThreadsPerWarp; + pipeline_compute_mma_ds_params.consumer_arv_count = 1; + pipeline_compute_mma_ds_params.initializing_warp = initializing_warp++; + PipelineComputeMmaDS pipeline_compute_mma_ds( + shared_storage.pipelines.compute_mma_ds, + pipeline_compute_mma_ds_params, + ClusterShape{}, + /*barrier init*/ cute::true_type{}, + /*mask calc*/ cute::false_type{}); + + typename PipelineMmaComputeDKDV::Params pipeline_mma_compute_dkdv_params; + if (role == WarpRole::Mma) { + pipeline_mma_compute_dkdv_params.role = + PipelineMmaComputeDKDV::ThreadCategory::Producer; + } + if (role == WarpRole::Compute) { + pipeline_mma_compute_dkdv_params.role = + PipelineMmaComputeDKDV::ThreadCategory::Consumer; + } + pipeline_mma_compute_dkdv_params.consumer_arv_count = + kNumComputeWarps * cutlass::NumThreadsPerWarp; + pipeline_mma_compute_dkdv_params.initializing_warp = initializing_warp++; + PipelineMmaComputeDKDV pipeline_mma_compute_dkdv( + shared_storage.pipelines.mma_compute_dkdv, + pipeline_mma_compute_dkdv_params, + ClusterShape{}, + /*barrier init*/ cute::true_type{}, + /*mask calc*/ cute::false_type{}); + PipelineReduceTmaStore pipeline_reduce_tma_store; + + TmemAllocator tmem_allocator; + + pipeline_init_arrive_relaxed(size(ClusterShape{})); + + pipeline_load_mma_q.init_masks(ClusterShape{}); + pipeline_load_mma_do.init_masks(ClusterShape{}); + pipeline_mma_compute_s.init_masks(ClusterShape{}); + pipeline_mma_compute_dp.init_masks(ClusterShape{}); + pipeline_mma_reduce_dq.init_masks(ClusterShape{}); + pipeline_compute_mma_p.init_masks(ClusterShape{}); + pipeline_compute_mma_ds.init_masks(ClusterShape{}); + pipeline_mma_compute_dkdv.init_masks(ClusterShape{}); + + typename decltype(pipeline_load_mma_q)::PipelineState + pipeline_load_mma_q_consumer_state; + typename decltype(pipeline_load_mma_do)::PipelineState + pipeline_load_mma_do_consumer_state; + typename decltype(pipeline_load_compute_lse)::PipelineState + pipeline_load_compute_lse_consumer_state; + typename decltype(pipeline_load_compute_sum_odo)::PipelineState + pipeline_load_compute_sum_odo_consumer_state; + typename decltype(pipeline_mma_compute_s)::PipelineState + pipeline_mma_compute_s_consumer_state; + typename decltype(pipeline_mma_compute_dp)::PipelineState + pipeline_mma_compute_dp_consumer_state; + typename decltype(pipeline_mma_reduce_dq)::PipelineState + pipeline_mma_reduce_dq_consumer_state; + typename decltype(pipeline_compute_mma_p)::PipelineState + pipeline_compute_mma_p_consumer_state; + typename decltype(pipeline_compute_mma_ds)::PipelineState + pipeline_compute_mma_ds_consumer_state; + typename decltype(pipeline_mma_compute_dkdv)::PipelineState + pipeline_mma_compute_dkdv_consumer_state; + + auto pipeline_load_mma_q_producer_state = + make_producer_start_state(); + auto pipeline_load_mma_do_producer_state = + make_producer_start_state(); + auto pipeline_load_compute_lse_producer_state = + make_producer_start_state(); + auto pipeline_load_compute_sum_odo_producer_state = + make_producer_start_state(); + auto pipeline_mma_compute_s_producer_state = + make_producer_start_state(); + auto pipeline_mma_compute_dp_producer_state = + make_producer_start_state(); + auto pipeline_mma_reduce_dq_producer_state = + make_producer_start_state(); + auto pipeline_compute_mma_p_producer_state = + make_producer_start_state(); + auto pipeline_compute_mma_ds_producer_state = + make_producer_start_state(); + auto pipeline_mma_compute_dkdv_producer_state = + make_producer_start_state(); + auto pipeline_reduce_tma_store_producer_state = + make_producer_start_state(); + + pipeline_init_wait(size(ClusterShape{})); + + auto blk_coord = make_coord( + _0{}, + blockIdx.x, + _0{}, + _0{}, + make_coord(make_coord(0, blockIdx.y), blockIdx.z)); + + // FNA's varlen mode works differently + auto blk_offset = make_coord( + _0{}, _0{}, _0{}, _0{}, make_coord(make_coord(_0{}, _0{}), _0{})); + auto problem_shape = params.problem_shape; + + // FNA + auto qkv_shape = params.fna.qkv_shape; + bool is_fully_block_sparse = params.fna.is_fully_block_sparse; + bool has_q_padding = params.fna.has_q_padding; + if (params.fna.requires_qkv_fixup) { + qkv_shape = Mask{}.correct_qkv_shape( + problem_shape, + params.fna.qkv_shape, + blk_coord, + params.fna.dilation, + params.fna.num_dilation_groups); + is_fully_block_sparse = fully_block_sparse( + qkv_shape, + get<0>(params.fna.na_params), + get<3>(params.fna.na_params), + QTileShape{}, + KVTileShape{}); + has_q_padding = not evenly_divides(qkv_shape, QTileShape{}); + } else if (params.fna.is_dilated) { + qkv_shape = ceil_div(params.fna.qkv_shape, params.fna.dilation); + is_fully_block_sparse = fully_block_sparse( + qkv_shape, + get<0>(params.fna.na_params), + get<3>(params.fna.na_params), + QTileShape{}, + KVTileShape{}); + has_q_padding = not evenly_divides(qkv_shape, QTileShape{}); + } + + auto [q_start, num_q_tiles] = Mask{}.get_trip_count( + blk_coord, + MultiDimTileShape{}, + params.fna.kv_shape, + qkv_shape, + params.fna.na_params); + + int iter_start = 0; + int iter_end = size(num_q_tiles); + int iter_count = (iter_end - iter_start) * get<4, 0, 0>(problem_shape); + + auto q_start_tile = ceil_div(q_start, QTileShape{}); + + auto q_tiled = ceil_div(params.fna.q_shape, QTileShape{}); + auto ctr = make_identity_tensor(num_q_tiles); + auto ctr_offset = domain_offset(q_start_tile, ctr); + + auto q_tiled_layout = make_layout(q_tiled); + + auto iter_to_tile_map = [&ctr_offset, &q_tiled_layout](int iter) { + return crd2idx(ctr_offset(iter), q_tiled_layout); + }; + + // Dilated cases where dilation doesn't evenly divide input shape can + // sometimes lead to threadblocks with no work to do (similar to varlen). We + // should exit those immediately instead of letting them proceed because + // some of the roles (i.e. load) unroll their loops, and this can cause + // stalls. + // End of FNA + if (iter_count <= 0) { + epilogue_clear( + blk_coord, + blk_offset, + problem_shape, + params.mainloop, + params.epilogue); + return; + } + + if (role == WarpRole::Load) { + warpgroup_reg_set(); + + load( + blk_coord, + blk_offset, + problem_shape, + iter_start, + iter_end, + iter_count, + iter_to_tile_map, + params.mainloop, + params.mainloop_params, + shared_storage.tensors, + pipeline_load_mma_q, + pipeline_load_mma_q_producer_state, + pipeline_load_mma_do, + pipeline_load_mma_do_producer_state, + pipeline_load_compute_lse, + pipeline_load_compute_lse_producer_state, + pipeline_load_compute_sum_odo, + pipeline_load_compute_sum_odo_producer_state); + + } else if (role == WarpRole::Mma) { + warpgroup_reg_set(); + + tmem_allocator.allocate( + TmemAllocator::Sm100TmemCapacityColumns, + &shared_storage.tmem_base_ptr); + __syncwarp(); + + mma(blk_coord, + problem_shape, + iter_start, + iter_end, + iter_count, + params.mainloop, + shared_storage.tensors, + pipeline_load_mma_q, + pipeline_load_mma_q_consumer_state, + pipeline_load_mma_do, + pipeline_load_mma_do_consumer_state, + pipeline_mma_compute_s, + pipeline_mma_compute_s_producer_state, + pipeline_mma_compute_dp, + pipeline_mma_compute_dp_producer_state, + pipeline_mma_reduce_dq, + pipeline_mma_reduce_dq_producer_state, + pipeline_compute_mma_p, + pipeline_compute_mma_p_consumer_state, + pipeline_compute_mma_ds, + pipeline_compute_mma_ds_consumer_state, + pipeline_mma_compute_dkdv, + pipeline_mma_compute_dkdv_producer_state); + + } else if (role == WarpRole::Compute) { + warpgroup_reg_set(); + + compute( + blk_coord, + blk_offset, + problem_shape, + iter_start, + iter_end, + iter_count, + // FNA params + params.fna.kv_shape, + qkv_shape, + params.fna.na_params, + q_start, + num_q_tiles, + is_fully_block_sparse, + has_q_padding, + // + params.mainloop, + params.epilogue, + shared_storage.tensors, + pipeline_load_compute_lse, + pipeline_load_compute_lse_consumer_state, + pipeline_load_compute_sum_odo, + pipeline_load_compute_sum_odo_consumer_state, + pipeline_mma_compute_s, + pipeline_mma_compute_s_consumer_state, + pipeline_mma_compute_dp, + pipeline_mma_compute_dp_consumer_state, + pipeline_compute_mma_p, + pipeline_compute_mma_p_producer_state, + pipeline_compute_mma_ds, + pipeline_compute_mma_ds_producer_state, + pipeline_mma_compute_dkdv, + pipeline_mma_compute_dkdv_consumer_state); + + cutlass::arch::NamedBarrier( + kNumComputeWarps * NumThreadsPerWarp, + cutlass::arch::ReservedNamedBarriers::EpilogueBarrier) + .arrive_and_wait(); + + if (warp_idx % kNumComputeWarps == 0) { + uint32_t free_stage_ptr = shared_storage.tmem_base_ptr; + tmem_allocator.free( + free_stage_ptr, TmemAllocator::Sm100TmemCapacityColumns); + } + + } else if (role == WarpRole::Reduce) { + warpgroup_reg_set(); + + reduce( + blk_coord, + problem_shape, + iter_start, + iter_end, + iter_count, + iter_to_tile_map, + params.mainloop, + params.mainloop_params, + shared_storage.tensors, + pipeline_mma_reduce_dq, + pipeline_mma_reduce_dq_consumer_state, + pipeline_reduce_tma_store, + pipeline_reduce_tma_store_producer_state); + + pipeline_reduce_tma_store.producer_tail( + pipeline_reduce_tma_store_producer_state); + } else { + warpgroup_reg_set(); + + /* no-op */ + } +#endif + } + + static dim3 get_block_shape() { + dim3 block(MaxThreadsPerBlock, 1, 1); + return block; + } + + static dim3 get_grid_shape(Params const& params) { + auto [Q, K, D, D_VO, HB] = params.problem_shape; + auto [H, B] = HB; + auto [H_R, H_K] = H; + dim3 grid(ceil_div(K, TileShapeK{}), H_K, B); + return grid; + } +}; + +} // namespace cutlass::fna::kernel diff --git a/natten/csrc/include/natten/cuda/fna_blackwell/kernel/sm100_fna_fwd_kernel_tma_warpspecialized.hpp b/natten/csrc/include/natten/cuda/fna_blackwell/kernel/sm100_fna_fwd_kernel_tma_warpspecialized.hpp new file mode 100644 index 00000000..57536837 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna_blackwell/kernel/sm100_fna_fwd_kernel_tma_warpspecialized.hpp @@ -0,0 +1,713 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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/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 "natten/cuda/fna_blackwell/collective/fna_common.hpp" +#include "natten/cuda/fna_blackwell/collective/fna_fusion.hpp" + +#include "natten/cuda/fmha_blackwell/kernel/fmha_tile_scheduler.hpp" + +namespace cutlass::fna::kernel { + +using namespace cute; +using namespace cutlass::fna::collective; + +struct Sm100FnaCtxKernelWarpspecializedSchedule { + 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 < + class ProblemShapeIn, + class CollectiveMainloop, + class CollectiveEpilogue, + class TileScheduler, + class KernelSchedule = Sm100FnaCtxKernelWarpspecializedSchedule> +struct Sm100FnaFwdKernelTmaWarpspecialized { + 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_assert(NumWarpsEpilogue == CollectiveEpilogue::NumWarpsEpilogue); + static_assert(NumWarpsLoad == CollectiveEpilogue::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; + + static constexpr bool IsMla = false; + + using ClusterShape = typename CollectiveMainloop::ClusterShape; + + using TmemAllocator = cute::TMEM::Allocator1Sm; + + struct SharedStorage { + using UnionType = union { + typename CollectiveMainloop::TensorStorage mainloop; + typename CollectiveEpilogue::TensorStorage epilogue; + }; + + using StructType = struct { + typename CollectiveMainloop::TensorStorage mainloop; + typename CollectiveEpilogue::TensorStorage epilogue; + }; + + static constexpr bool IsPersistent = std::is_same_v< + TileScheduler, + cutlass::fmha::kernel::PersistentTileScheduler>; + using MainloopEpilogueStorage = std::conditional_t< + IsPersistent, + std::conditional_t< + IsMla, + std::conditional_t< + CollectiveMainloop::IsOrderLoadEpilogue, + UnionType, + StructType>, + StructType>, + UnionType>; + + MainloopEpilogueStorage mainloop_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); + + struct Arguments { + ProblemShape problem_shape; + typename CollectiveMainloop::Arguments mainloop; + typename CollectiveEpilogue::Arguments epilogue; + 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.problem_shape, args.hw_info, ClusterShape{}, TileShape{})}; + } + + CUTLASS_DEVICE auto apply_batch( + const Params& params, + ProblemShape const& problem_shape, + int batch_idx) { + return params.problem_shape; + } + + CUTLASS_DEVICE void operator()(const Params& params, char* smem) { +#if ( \ + !defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) && \ + !defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) && \ + !defined(CUTLASS_ARCH_MMA_SM103A_ENABLED) && \ + !defined(CUTLASS_ARCH_MMA_SM103F_ENABLED)) + printf( + "ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); +#else + + 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); + + auto get_epilogue_storage = [&]() { + if constexpr (IsMla && CollectiveMainloop::IsOrderLoadEpilogue) { + return reinterpret_cast( + shared_storage.mainloop_epilogue.mainloop.smem_o.data()); + } else { + return &shared_storage.mainloop_epilogue.epilogue; + } + }; + typename CollectiveEpilogue::TensorStorage& epilogue_storage = + *get_epilogue_storage(); + + 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::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.is_leader = + lane_predicate && (role == WarpRole::Load); + pipeline_load_kv_params.transaction_bytes = + CollectiveMainloop::TransactionBytesLoadK; + 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::PipelineQ>(); + + 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::PipelineKV>(); + + 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>(); + + 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::PipelineS>(); + + 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>(); + + 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::PipelineC>(); + + 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::PipelineE>(); + + typename CollectiveMainloop::PipelineO::PipelineState + pipeline_mma_corr_consumer_state; + typename CollectiveMainloop::PipelineO::PipelineState + pipeline_mma_corr_producer_state = cutlass::make_producer_start_state< + typename CollectiveMainloop::PipelineO>(); + + 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) { + 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(); + + bool has_valid = false; + + 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; + } + + has_valid = true; + + if (get<1>(logical_problem_shape) == 0) { + mainloop.correction_empty( + blk_coord, + params.mainloop, + logical_problem_shape, + params.problem_shape, + epilogue_storage, + pipeline_corr_epi, + pipeline_corr_epi_producer_state, + epilogue); + continue; + } + + mainloop.correction( + blk_coord, + params.mainloop, + logical_problem_shape, + epilogue_storage, + 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); + + if (has_valid) { + 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(); + + bool allocated = false; + + 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 (!allocated) { + tmem_allocator.allocate( + TmemAllocator::Sm100TmemCapacityColumns, + &shared_storage.tmem_base_ptr); + __syncwarp(); + allocated = true; + } + + if (get<1>(logical_problem_shape) == 0) { + continue; + } + + mainloop.mma( + blk_coord, + params.mainloop, + logical_problem_shape, + shared_storage.mainloop_epilogue.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(); + + if constexpr (IsMla && CollectiveMainloop::IsOrderLoadEpilogue) { + cutlass::arch::NamedBarrier::arrive( + (NumWarpsLoad + NumWarpsEpilogue) * NumThreadsPerWarp, + cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); + } + + 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) { + continue; + } + + mainloop.load( + blk_coord, + logical_problem_shape, + params.mainloop, + params.problem_shape, + shared_storage.mainloop_epilogue.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(); + + bool has_valid = false; + + 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; + } + + has_valid = true; + + epilogue.store( + blk_coord, + logical_problem_shape, + params.epilogue, + params.problem_shape, + epilogue_storage, + pipeline_corr_epi, + pipeline_corr_epi_consumer_state); + } + + static_assert(NumWarpsEpilogue <= 1); + if constexpr (NumWarpsEpilogue == 1) { + if (has_valid) { + 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 */ + } +#endif + } +}; + +} // namespace cutlass::fna::kernel diff --git a/natten/csrc/include/natten/cuda/fna_hopper/collective/fna_collective_bwd_tma_warpspecialized.hpp b/natten/csrc/include/natten/cuda/fna_hopper/collective/fna_collective_bwd_tma_warpspecialized.hpp new file mode 100644 index 00000000..4e7e84bf --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna_hopper/collective/fna_collective_bwd_tma_warpspecialized.hpp @@ -0,0 +1,1422 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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/gemm/collective/collective_builder.hpp" + +#include "natten/cuda/fna_hopper/collective/fna_collective_load.hpp" +#include "natten/cuda/fna_hopper/collective/fna_common.hpp" +#include "natten/cuda/fna_hopper/collective/fna_fusion.hpp" +#include "natten/cuda/fna_hopper/collective/fna_fusion_bwd.hpp" + +// This is identical to FMHA; reuse +#include "natten/cuda/fmha_hopper/kernel/fmha_options.hpp" + +namespace cutlass::fna::collective { + +using cutlass::fmha::kernel::find_option_t; +using cutlass::fmha::kernel::Tag; + +template < + typename Element_, + typename ElementAccumulator_, + typename TileShape_, // BlockQO, BlockKV, BlockHead + class Mask, + class QTileShape, + class KVTileShape, + class NADim, + class... Options> +struct FnaBwdMainloopTmaWarpSpecializedSm90 { + using MultiDimTileShape = cute::tuple; + + static_assert( + size(QTileShape{}) == get<0>(TileShape_{}), + "QTileShape doesn't match the size of Q tile in the FMHA kernel."); + static_assert( + size(KVTileShape{}) == get<1>(TileShape_{}), + "KVTileShape doesn't match the size of KV tile in the FMHA kernel."); + + using Element = Element_; + using ElementAccumulator = ElementAccumulator_; + using TileShape = TileShape_; + + static constexpr bool kIsPersistent = false; + + static const int NumLoadWarpGroups = 1; + static constexpr int NumMmaWarpGroups = 2; + static constexpr int StageCountQ = 2 /*K, V*/ * NumMmaWarpGroups; + static constexpr int StageCount = 2 /*Q, dO*/ * 2 /* actual stages */; + + static const int kOuterLoads = 2; + using StagesQ = cutlass::gemm::collective::StageCount; + using Stages = cutlass::gemm::collective::StageCount; + using ClusterShape = Shape<_1, _1, _1>; + static_assert(StagesQ::value >= 2); + static_assert(Stages::value >= 2 * NumMmaWarpGroups); + + // 16B alignment lets us use TMA + static constexpr int Alignment = 16 / sizeof(Element); + + using TileShapeNM = Shape< // (N,M,D) + decltype(tuple_element_t<1, TileShape>{} / Int{}), + tuple_element_t<0, TileShape>, + tuple_element_t<2, TileShape>>; + + using TileShapeND = decltype(select<0, 2, 1>(TileShapeNM{})); // (N,D,M) + + using TileShapeMD = decltype(select<2, 1, 0>(TileShapeND{})); // (M,D,N) + + using CollectiveMmaNM = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm90, + cutlass::arch::OpClassTensorOp, + Element, + cute::tuple>, + Alignment, + Element, + cute::tuple>, + Alignment, + ElementAccumulator, + TileShapeNM, + ClusterShape, + Stages, + cutlass::gemm::KernelTmaWarpSpecialized>::CollectiveOp; + + using CollectiveMmaND = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm90, + cutlass::arch::OpClassTensorOp, + Element, + cute::tuple>, + Alignment, // from register, doesn't matter + Element, + cute::tuple<_1, int, cute::tuple>, + Alignment, + ElementAccumulator, + TileShapeND, + ClusterShape, + Stages, + cutlass::gemm::KernelTmaWarpSpecialized>::CollectiveOp; + + using CollectiveMmaND_SS = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm90, + cutlass::arch::OpClassTensorOp, + Element, + cute::tuple>, + Alignment, // from register, doesn't matter + Element, + cute::tuple<_1, int, cute::tuple>, + Alignment, + ElementAccumulator, + TileShapeND, + ClusterShape, + Stages, + cutlass::gemm::KernelTmaWarpSpecialized>::CollectiveOp; + + using CollectiveMmaMD = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm90, + cutlass::arch::OpClassTensorOp, + Element, + cute::tuple<_1, int, cute::tuple>, + Alignment, // from smem, might matter (?) + Element, + cute::tuple<_1, int, cute::tuple>, + Alignment, + ElementAccumulator, + TileShapeMD, + ClusterShape, + Stages, + cutlass::gemm::KernelTmaWarpSpecialized>::CollectiveOp; + + using TiledMmaNM = typename CollectiveMmaNM::TiledMma; + using TiledMmaND_SS = typename CollectiveMmaND_SS::TiledMma; + using TiledMmaND_RS = + decltype(convert_to_gmma_rs(typename CollectiveMmaND::TiledMma{})); + using TiledMmaND = TiledMmaND_RS; + using TiledMmaMD = typename CollectiveMmaMD::TiledMma; + + using SmemLayoutQ = typename CollectiveMmaNM::SmemLayoutB; + using SmemLayoutK = typename CollectiveMmaNM::SmemLayoutA; + using SmemLayoutV = typename CollectiveMmaNM::SmemLayoutA; + using SmemLayoutDO = typename CollectiveMmaNM::SmemLayoutB; + + // using SmemLayoutDQ = Layout< + // Shape< + // tuple_element_t<0, TileShapeMD>, + // Shape<_2, _4, decltype(tuple_element_t<1, TileShapeMD>{} / _8{})>, + // _2 + // >, + // Stride< + // _4, + // Stride{} * _4{}), _1, + // decltype(tuple_element_t<0, TileShapeMD>{} * _8{})>, + // decltype(tuple_element_t<0, TileShapeMD>{} * tuple_element_t<1, + // TileShapeMD>{}) + // >>; + + using SmemLayoutDQ_0 = Layout< + Shape< + tuple_element_t<0, TileShapeMD>, + tuple_element_t<1, TileShapeMD>, + _2>, + Stride< + tuple_element_t<1, TileShapeMD>, + _1, + decltype(tuple_element_t<0, TileShapeMD>{} * tuple_element_t<1, TileShapeMD>{})>>; + + using SmemAtomDQ = + decltype(cutlass::gemm::collective::detail::ss_smem_selector< + cute::GMMA::Major::K, + ElementAccumulator, + tuple_element_t<0, TileShapeMD>, + tuple_element_t<1, TileShapeMD>>()); + using SmemLayoutDQ_1 = decltype(tile_to_shape( + SmemAtomDQ{}, + make_shape(get<0>(TileShapeMD{}), get<1>(TileShapeMD{}), _2{}), + Step<_2, _1, _3>{})); + using SmemLayoutDQ = SmemLayoutDQ_1; + + using PipelineDQ = cutlass::PipelineAsync<2>; + + using SmemLayoutDS_0 = decltype(unstageSmemLayout( + typename CollectiveMmaMD::SmemLayoutA{}, + Int{})); + + using SmemLayoutDS = decltype(tile_to_shape( + GMMA::Layout_MN_INTER_Atom{}, + make_shape( + size<0>(SmemLayoutDS_0{}), + size<1>(SmemLayoutDS_0{}), + size<2>(SmemLayoutDS_0{})), + Step<_1, _2, _3>{})); + using SmemLayoutKp = typename CollectiveMmaMD::SmemLayoutB; + + using SmemLayoutQp = typename CollectiveMmaND::SmemLayoutB; + using SmemLayoutDOp = typename CollectiveMmaND::SmemLayoutB; + + using SmemLayoutLSE = + Layout, Int>>; + + using MainloopPipeline = cutlass::PipelineTmaAsync; + using MainloopPipelineQ = cutlass::PipelineTmaAsync; + + using PipelineState = + typename cutlass::PipelineState; + using PipelineStateQ = + typename cutlass::PipelineState; + + using TileShapePV = TileShapeND; // To work with the kernel level + using TiledMmaPV = TiledMmaND; + + static constexpr int kInnerLoadBytes = + size(SmemLayoutQ{}(_, _, _0{})) * sizeof(Element) + + size(SmemLayoutLSE{}(_, _0{})) * sizeof(ElementAccumulator); + static constexpr int kOuterLoadBytes = + size(SmemLayoutK{}(_, _, _0{})) * sizeof(Element); + + struct SharedStorage { + // One for each consumer WG + union { + cute::array_aligned> smem_k; + cute::array_aligned> smem_kp; + cute::array_aligned> smem_v; + }; + + cute::array_aligned> smem_ds; + + // Loaded by producer, consumed by both WGs + union { + cute::array_aligned> smem_q; + cute::array_aligned> smem_do; + cute::array_aligned> smem_qp; + cute::array_aligned> smem_dop; + }; + + // Accumulated into by both consumers, potentially loaded, potentially + // written + cute::array_aligned> + smem_dq; + + union { + cute::array_aligned> + smem_lse; + cute::array_aligned> + smem_sumOdO; + }; + }; + + struct FnaArguments { + NADim q_shape; + NADim kv_shape; + NADim qkv_shape; + NADim window_size; + NADim stride; + NADim dilation; + }; + + struct Arguments { + const Element* ptr_Q; + cute::tuple dQ; + const Element* ptr_K; + cute::tuple dK; + const Element* ptr_V; + cute::tuple dV; + + const Element* ptr_dO; + cute::tuple dDO; + + const ElementAccumulator* ptr_LSE; + cute::tuple dLSE; + const ElementAccumulator* ptr_sum_OdO; + cute::tuple dSumOdO; + + ElementAccumulator* ptr_dQ; + cute::tuple dDQ; + + // if zero, defaults to 1/sqrt(D) + float scale_softmax = 0.0f; + + FnaArguments fna; + }; + + using TMA_Q = typename CollectiveMmaNM::Params::TMA_B; + using TMA_K = typename CollectiveMmaNM::Params::TMA_A; + using TMA_V = typename CollectiveMmaNM::Params::TMA_A; + using TMA_DO = typename CollectiveMmaNM::Params::TMA_B; + + using TMA_LSE = decltype(make_tma_copy( + SM90_TMA_LOAD{}, + make_tensor( + (const ElementAccumulator*)nullptr, + make_shape(1, 1, 1), + make_stride(_1{}, 0, 0)), + SmemLayoutLSE{}(_, _0{}))); + using TMA_ODO = TMA_LSE; + + using TMA_DQ = decltype(make_tma_copy( + SM90_TMA_REDUCE_ADD{}, + make_tensor( + (const ElementAccumulator*)nullptr, + make_shape(1, 1, 1, 1), + make_stride(0, _1{}, 0, 0)), + SmemLayoutDQ{}(_, _, _0{}))); + + using LoadQ = CollectiveLoadTma< + LoadKind::kBwdM, + MainloopPipeline, + Element, + SmemLayoutQ, + TMA_Q>; + + using LoadK = CollectiveLoadTma< + LoadKind::kBwdN, + MainloopPipelineQ, + Element, + SmemLayoutK, + TMA_K>; + + using LoadV = CollectiveLoadTma< + LoadKind::kBwdN, + MainloopPipelineQ, + Element, + SmemLayoutV, + TMA_V>; + + using LoadDO = CollectiveLoadTma< + LoadKind::kBwdM, + MainloopPipeline, + Element, + SmemLayoutDO, + TMA_DO>; + + using LoadLSE = CollectiveLoadTma< + LoadKind::kBwdScalar, + MainloopPipeline, + ElementAccumulator, + SmemLayoutLSE, + TMA_LSE>; + + using LoadODO = CollectiveLoadTma< + LoadKind::kBwdScalar, + MainloopPipeline, + ElementAccumulator, + SmemLayoutLSE, + TMA_ODO>; + + struct FnaParams { + NADim qkv_shape; + NADim q_shape; + NADim kv_shape; + cute::tuple + na_params; // win, win_left, win_right, stride, stride_offset + bool is_fully_block_sparse; + bool has_q_padding; + NADim dilation; + bool requires_qkv_fixup; + bool is_dilated; + int num_dilation_groups; + }; + + struct Params { + TMA_Q tma_load_q; + TMA_K tma_load_k; + TMA_V tma_load_v; + TMA_DO tma_load_do; + + TMA_LSE tma_load_lse; + TMA_ODO tma_load_odo; + + TMA_DQ tma_red_dq; + + float scale_softmax; + float scale_softmax_log2; + + FnaParams fna; + }; + + static_assert(size(TiledMmaNM{}) == size(TiledMmaND{})); + static_assert(size(TiledMmaNM{}) == size(TiledMmaMD{})); + + template + static bool can_implement( + ProblemShape const& problem_size, + Arguments const& args) { + // NOTE (ali): the original Hopper FMHA kernel does an alignment check on + // seqlen_Q as well, citing TMA requirement. However, it does not look like + // TMA runs into any issues, because cases that violate this run perfectly + // fine. + // + // Backprop requires this extra check though, and the point of failure is + // LSE. + // + // This does not cause any issues with FNA kernels, since they already pad + // for tile size, and tile sizes are usually powers of two, and at least 64, + // so divisibility by the required alignment 16 is a guarantee, though not + // explicitly asserted. + // + // Adding in padding along seqlen_Q requires modifications in the bwd mask + // as well. + // + // Therefore, we only pad LSE, and not any other tensors to get around this + // issue in bwd, instead of padding Q, dQ, O and dO as well. + // + // If we pad those, we would have to modify the residual mask as well. + // + return true && (get<4>(problem_size) <= get<2>(TileShape{})) && + ((get<4>(problem_size) % Alignment) == 0) && + /*((get<2>(problem_size) % Alignment) == 0)*/ + evenly_divides(args.fna.q_shape, QTileShape{}) && + evenly_divides(args.fna.kv_shape, KVTileShape{}) && + evenly_divides( + get<0>(problem_size), + size(args.fna.dilation)) && // dilation groups are + // folded into batch + tuple_leq(args.fna.window_size, args.fna.qkv_shape) && + tuple_leq(args.fna.stride, args.fna.window_size); + } + + template + static Params to_underlying_arguments( + ProblemShape const& problem_size, + Arguments const& args, + void* workspace) { + auto problem_shape_nm = make_shape( + get<3>(problem_size), + get<2>(problem_size), + get<4>(problem_size), + make_shape(get<0>(problem_size), get<1>(problem_size))); + + auto dK = make_stride( + get<2>(args.dK), + get<3>(args.dK), + make_stride(get<0>(args.dK), get<1>(args.dK))); + auto dQ = make_stride( + get<2>(args.dQ), + get<3>(args.dQ), + make_stride(get<0>(args.dQ), get<1>(args.dQ))); + auto params_nm_kq = CollectiveMmaNM::to_underlying_arguments( + problem_shape_nm, + typename CollectiveMmaNM::Arguments{ + args.ptr_K, + dK, + args.ptr_Q, + dQ, + }, + /*workspace=*/nullptr); + + auto dV = make_stride( + get<2>(args.dV), + get<3>(args.dV), + make_stride(get<0>(args.dV), get<1>(args.dV))); + auto dDO = make_stride( + get<2>(args.dDO), + get<3>(args.dDO), + make_stride(get<0>(args.dDO), get<1>(args.dDO))); + auto params_nm_vdo = CollectiveMmaNM::to_underlying_arguments( + problem_shape_nm, + typename CollectiveMmaNM::Arguments{ + args.ptr_V, + dV, + args.ptr_dO, + dDO, + }, + /*workspace=*/nullptr); + + TMA_LSE tma_load_lse = make_tma_copy( + SM90_TMA_LOAD{}, + make_tensor( + args.ptr_LSE, + select<2, 0, 1>(problem_size), + select<2, 0, 1>(args.dLSE)), + SmemLayoutLSE{}(_, _0{})); + TMA_ODO tma_load_odo = make_tma_copy( + SM90_TMA_LOAD{}, + make_tensor( + args.ptr_sum_OdO, + select<2, 0, 1>(problem_size), + select<2, 0, 1>(args.dSumOdO)), + SmemLayoutLSE{}(_, _0{})); + + TMA_DQ tma_red_dq = make_tma_copy( + SM90_TMA_REDUCE_ADD{}, + make_tensor( + args.ptr_dQ, + select<2, 4, 0, 1>(problem_size), + select<2, 3, 0, 1>(args.dDQ)), + SmemLayoutDQ{}(_, _, _0{})); + + float scale_softmax = args.scale_softmax; + if (scale_softmax == 0.0f) { + scale_softmax = 1.0f / (float)std::sqrt(get<4>(problem_size)); + } + float log2_e = static_cast(std::log2(std::exp(1.0))); + + // FNA params + auto window_left = get_window_left(args.fna.window_size); + auto window_right = get_window_right(args.fna.window_size); + auto stride_offset = get_bwd_stride_offset(args.fna.stride); + auto na_params = make_tuple( + args.fna.window_size, + window_left, + window_right, + args.fna.stride, + stride_offset); + bool requires_qkv_fixup = + not evenly_divides(args.fna.qkv_shape, args.fna.dilation); + + auto is_fully_block_sparse = fully_block_sparse( + args.fna.qkv_shape, + args.fna.window_size, + args.fna.stride, + QTileShape{}, + KVTileShape{}); + + return Params{ + params_nm_kq.tma_load_b, + params_nm_kq.tma_load_a, + params_nm_vdo.tma_load_a, + params_nm_vdo.tma_load_b, + tma_load_lse, + tma_load_odo, + tma_red_dq, + scale_softmax, + log2_e * scale_softmax, + FnaParams{ + args.fna.qkv_shape, + args.fna.q_shape, + args.fna.kv_shape, + na_params, + is_fully_block_sparse, + /* has_q_padding */ + not evenly_divides(args.fna.qkv_shape, QTileShape{}), + args.fna.dilation, + requires_qkv_fixup, + is_dilated(args.fna.dilation), + size(args.fna.dilation)}, + }; + } + + template + CUTLASS_DEVICE auto get_trip_info( + Params const& params, + BlkCoord const& blk_coord, + ProblemShape const& problem_shape) { + auto qkv_shape = params.fna.qkv_shape; + bool is_fully_block_sparse = params.fna.is_fully_block_sparse; + bool has_q_padding = params.fna.has_q_padding; + if (params.fna.requires_qkv_fixup) { + qkv_shape = Mask{}.correct_qkv_shape( + problem_shape, + params.fna.qkv_shape, + blk_coord, + params.fna.dilation, + params.fna.num_dilation_groups); + is_fully_block_sparse = fully_block_sparse( + qkv_shape, + get<0>(params.fna.na_params), + get<3>(params.fna.na_params), + QTileShape{}, + KVTileShape{}); + has_q_padding = not evenly_divides(qkv_shape, QTileShape{}); + } else if (params.fna.is_dilated) { + qkv_shape = ceil_div(params.fna.qkv_shape, params.fna.dilation); + is_fully_block_sparse = fully_block_sparse( + qkv_shape, + get<0>(params.fna.na_params), + get<3>(params.fna.na_params), + QTileShape{}, + KVTileShape{}); + has_q_padding = not evenly_divides(qkv_shape, QTileShape{}); + } + + auto [q_start, num_q_tiles] = Mask{}.get_trip_count( + blk_coord, + MultiDimTileShape{}, + params.fna.kv_shape, + qkv_shape, + params.fna.na_params); + + return cute::make_tuple( + q_start, num_q_tiles, qkv_shape, is_fully_block_sparse, has_q_padding); + } + + 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()); + cute::prefetch_tma_descriptor(params.tma_load_do.get_tma_descriptor()); + cute::prefetch_tma_descriptor(params.tma_load_odo.get_tma_descriptor()); + cute::prefetch_tma_descriptor(params.tma_load_lse.get_tma_descriptor()); + } + + template < + bool kLoadOuter, + class BlkCoord, + class ProblemShape, + class LoadWarpBarrier> + CUTLASS_DEVICE void load_kv_maybe_q( + int block_rank_in_cluster, + BlkCoord const& blk_coord, + Params const& params, + ProblemShape const& problem_size, + MainloopPipeline& pipeline_inner, + PipelineState& smem_pipe_write_inner, + MainloopPipelineQ& pipeline_outer, + PipelineStateQ& smem_pipe_write_outer, + SharedStorage& storage, + LoadWarpBarrier& load_warp_barrier, + bool do_barrier) { + // Load pattern: + // K0 V0 K1 V1 + // Q0 DO0 Q1 DO1 Q2 DO2 ... + // K0 Q0 V0 K1 DO0 V1 ... + int lane_predicate = cute::elect_one_sync(); + + int outer_tile_count = NumMmaWarpGroups; + + auto + [q_start, + num_q_tiles, + qkv_shape, + is_fully_block_sparse, + has_q_padding] = get_trip_info(params, blk_coord, problem_size); + int inner_tile_count = size(num_q_tiles); + + auto q_start_tile = ceil_div(q_start, QTileShape{}); + + auto q_tiled = ceil_div(params.fna.q_shape, QTileShape{}); + auto ctr = make_identity_tensor(num_q_tiles); + auto ctr_offset = domain_offset(q_start_tile, ctr); + + auto q_tiled_layout = make_layout(q_tiled); + + auto iter_to_tile_map = [&ctr_offset, &q_tiled_layout](int iter) { + return crd2idx(ctr_offset(iter), q_tiled_layout); + }; + + // Dilated cases where dilation doesn't evenly divide input shape can + // sometimes lead to threadblocks with no work to do (similar to varlen). We + // should exit those immediately instead of letting them proceed because + // some of the roles (i.e. load) unroll their loops, and this can cause + // stalls. + if (inner_tile_count == 0) { + return; + } + + auto outer_tile_iter = cute::make_coord_iterator(outer_tile_count); + auto inner_tile_iter = cute::make_coord_iterator(inner_tile_count); + + uint16_t mcast_mask_b = 0; + + LoadQ load_q{params.tma_load_q, pipeline_inner, storage.smem_q}; + auto load_state_q = load_q.init_state( + block_rank_in_cluster, + problem_size, + TileShapeNM{}, + blk_coord, + inner_tile_count); + + LoadDO load_do{params.tma_load_do, pipeline_inner, storage.smem_do}; + auto load_state_do = load_do.init_state( + block_rank_in_cluster, + problem_size, + TileShapeNM{}, + blk_coord, + inner_tile_count); + + LoadK load_k{params.tma_load_k, pipeline_outer, storage.smem_k}; + auto load_state_k = load_k.init_state( + _0{}, problem_size, TileShapeNM{}, blk_coord, outer_tile_count); + + LoadV load_v{params.tma_load_v, pipeline_outer, storage.smem_v}; + auto load_state_v = load_v.init_state( + _0{}, problem_size, TileShapeNM{}, blk_coord, outer_tile_count); + + LoadLSE load_lse{params.tma_load_lse, pipeline_inner, storage.smem_lse}; + auto load_state_lse = load_lse.init_state( + _0{}, problem_size, TileShapeNM{}, blk_coord, outer_tile_count); + + LoadODO load_odo{params.tma_load_odo, pipeline_inner, storage.smem_sumOdO}; + auto load_state_odo = load_odo.init_state( + _0{}, problem_size, TileShapeNM{}, blk_coord, outer_tile_count); + + outer_tile_count *= 2; // K & V + inner_tile_count *= 4; // Q & dO & LSE & sumOdO + + if constexpr (kLoadOuter) { + load_k.template step( + outer_tile_iter, + load_state_k, + smem_pipe_write_outer, + lane_predicate, + outer_tile_count); + } + + load_q.template step_with_iter_tile_map( + inner_tile_iter, + load_state_q, + smem_pipe_write_inner, + lane_predicate, + inner_tile_count, + mcast_mask_b, + iter_to_tile_map); + load_lse.template step_with_iter_tile_map( + inner_tile_iter, + load_state_lse, + smem_pipe_write_inner, + lane_predicate, + inner_tile_count, + mcast_mask_b, + iter_to_tile_map); + + if constexpr (!kLoadOuter) { + if (do_barrier) { + load_warp_barrier.arrive(); + load_warp_barrier.wait(/*phase=*/0); + do_barrier = false; + } + } + + if constexpr (kLoadOuter) { + load_v.template step( + outer_tile_iter, + load_state_v, + smem_pipe_write_outer, + lane_predicate, + outer_tile_count); + load_k.template step( + outer_tile_iter, + load_state_k, + smem_pipe_write_outer, + lane_predicate, + outer_tile_count); + } + + load_do.template step_with_iter_tile_map( + inner_tile_iter, + load_state_do, + smem_pipe_write_inner, + lane_predicate, + inner_tile_count, + mcast_mask_b, + iter_to_tile_map); + load_odo.template step_with_iter_tile_map( + inner_tile_iter, + load_state_odo, + smem_pipe_write_inner, + lane_predicate, + inner_tile_count, + mcast_mask_b, + iter_to_tile_map); + + if constexpr (kLoadOuter) { + load_v.template step( + outer_tile_iter, + load_state_v, + smem_pipe_write_outer, + lane_predicate, + outer_tile_count); + } + + if constexpr (kLoadOuter) { + while (outer_tile_count > 0) { + load_k.template step( + outer_tile_iter, + load_state_k, + smem_pipe_write_outer, + lane_predicate, + outer_tile_count); + load_v.template step( + outer_tile_iter, + load_state_v, + smem_pipe_write_outer, + lane_predicate, + outer_tile_count); + } + } + + CUTLASS_PRAGMA_NO_UNROLL + while (inner_tile_count > 0) { + load_q.template step_with_iter_tile_map( + inner_tile_iter, + load_state_q, + smem_pipe_write_inner, + lane_predicate, + inner_tile_count, + mcast_mask_b, + iter_to_tile_map); + load_lse.template step_with_iter_tile_map( + inner_tile_iter, + load_state_lse, + smem_pipe_write_inner, + lane_predicate, + inner_tile_count, + mcast_mask_b, + iter_to_tile_map); + + load_do.template step_with_iter_tile_map( + inner_tile_iter, + load_state_do, + smem_pipe_write_inner, + lane_predicate, + inner_tile_count, + mcast_mask_b, + iter_to_tile_map); + load_odo.template step_with_iter_tile_map( + inner_tile_iter, + load_state_odo, + smem_pipe_write_inner, + lane_predicate, + inner_tile_count, + mcast_mask_b, + iter_to_tile_map); + } + } + + template + CUTLASS_DEVICE void load_maybe_q( + BlkCoord const& blk_coord, + Params const& params, + ProblemShape const& problem_size, + MainloopPipelineQ& pipeline_outer, + PipelineStateQ& smem_pipe_write_outer, + SharedStorage& storage, + LoadWarpBarrier& load_warp_barrier, + bool do_barrier) { + auto + [q_start, + num_q_tiles, + qkv_shape, + is_fully_block_sparse, + has_q_padding] = get_trip_info(params, blk_coord, problem_size); + int inner_tile_count = size(num_q_tiles); + + // Dilated cases where dilation doesn't evenly divide input shape can + // sometimes lead to threadblocks with no work to do (similar to varlen). We + // should exit those immediately instead of letting them proceed because + // some of the roles (i.e. load) unroll their loops, and this can cause + // stalls. + if (inner_tile_count == 0) { + return; + } + + // Load pattern: + // K0 V0 K1 V1 + // Q0 DO0 Q1 DO1 Q2 DO2 ... + // K0 Q0 V0 K1 DO0 V1 ... + int lane_predicate = cute::elect_one_sync(); + + int outer_tile_count = NumMmaWarpGroups; + + auto outer_tile_iter = cute::make_coord_iterator(outer_tile_count); + + LoadK load_k{params.tma_load_k, pipeline_outer, storage.smem_k}; + auto load_state_k = load_k.init_state( + _0{}, problem_size, TileShapeNM{}, blk_coord, outer_tile_count); + + LoadV load_v{params.tma_load_v, pipeline_outer, storage.smem_v}; + auto load_state_v = load_v.init_state( + _0{}, problem_size, TileShapeNM{}, blk_coord, outer_tile_count); + + outer_tile_count *= 2; // K & V + + load_k.template step( + outer_tile_iter, + load_state_k, + smem_pipe_write_outer, + lane_predicate, + outer_tile_count); + + if (do_barrier) { + load_warp_barrier.arrive(); + load_warp_barrier.wait(/*phase=*/0); + do_barrier = false; + } + + load_v.template step( + outer_tile_iter, + load_state_v, + smem_pipe_write_outer, + lane_predicate, + outer_tile_count); + + while (outer_tile_count > 0) { + load_k.template step( + outer_tile_iter, + load_state_k, + smem_pipe_write_outer, + lane_predicate, + outer_tile_count); + load_v.template step( + outer_tile_iter, + load_state_v, + smem_pipe_write_outer, + lane_predicate, + outer_tile_count); + } + } + + template < + class BlkCoord, + class ProblemShape, + class MainloopPipelineReducer, + class PipelineStateReducer> + CUTLASS_DEVICE void reduce( + BlkCoord const& blk_coord, + Params const& params, + ProblemShape const& problem_size, + MainloopPipelineReducer& pipeline_reducer, + PipelineStateReducer& smem_pipe_read_reducer, + SharedStorage& storage) { + int lane_predicate = cute::elect_one_sync(); + + Tensor mDQ_full = + params.tma_red_dq.get_tma_tensor(select<2, 4, 0, 1>(problem_size)); + Tensor gDQ_full = local_tile( + mDQ_full, + TileShapeMD{}, + make_coord(_, _, _), + Step<_1, _1, Underscore>{}); + Tensor gDQ = + gDQ_full(_, _, _, _0{}, get<2, 0>(blk_coord), get<2, 1>(blk_coord)); + Tensor sDQ = + make_tensor(make_smem_ptr(storage.smem_dq.data()), SmemLayoutDQ{}); + + auto block_tma = params.tma_red_dq.get_slice(_0{}); + + Tensor tDQsDQ = block_tma.partition_S(sDQ); + Tensor tDQgDQ = block_tma.partition_D(gDQ); + + auto + [q_start, + num_q_tiles, + qkv_shape, + is_fully_block_sparse, + has_q_padding] = get_trip_info(params, blk_coord, problem_size); + int inner_tile_count = size(num_q_tiles); + + auto q_start_tile = ceil_div(q_start, QTileShape{}); + + auto q_tiled = ceil_div(params.fna.q_shape, QTileShape{}); + auto ctr = make_identity_tensor(num_q_tiles); + auto ctr_offset = domain_offset(q_start_tile, ctr); + + auto q_tiled_layout = make_layout(q_tiled); + + auto iter_to_tile_map = [&ctr_offset, &q_tiled_layout](int iter) { + return crd2idx(ctr_offset(iter), q_tiled_layout); + }; + + // Dilated cases where dilation doesn't evenly divide input shape can + // sometimes lead to threadblocks with no work to do (similar to varlen). We + // should exit those immediately instead of letting them proceed because + // some of the roles (i.e. load) unroll their loops, and this can cause + // stalls. + if (inner_tile_count == 0) { + return; + } + + int g_index = 0; + + auto smem_pipe_release_reducer = smem_pipe_read_reducer; + bool first = true; + while (inner_tile_count > 0) { + if (inner_tile_count == 0) + break; + + pipeline_reducer.consumer_wait(smem_pipe_read_reducer); + if (lane_predicate == 1) { + tma_store_wait<1>(); + } + if (!first) { + pipeline_reducer.consumer_release(smem_pipe_release_reducer); + ++smem_pipe_release_reducer; + } else { + first = false; + } + if (lane_predicate == 1) { + copy( + params.tma_red_dq, + tDQsDQ(_, _, _, smem_pipe_read_reducer.index()), + tDQgDQ(_, _, _, iter_to_tile_map(g_index))); + tma_store_arrive(); + } + ++smem_pipe_read_reducer; + --inner_tile_count; + ++g_index; + } + if (lane_predicate) { + tma_store_wait<0>(); + } + pipeline_reducer.consumer_release(smem_pipe_release_reducer); + ++smem_pipe_release_reducer; + } + + template < + class BlkCoord, + class ProblemShape, + class MainloopPipelineReducer, + class PipelineStateReducer, + class MathWgOrderBarrier> + CUTLASS_DEVICE auto compute( + BlkCoord const& blk_coord, + BlkCoord const& wg_coord, + Params const& params, + ProblemShape const& problem_size, + MainloopPipeline& pipeline_inner, + PipelineState& smem_pipe_read_inner, + MainloopPipelineQ& pipeline_outer, + PipelineStateQ& smem_pipe_read_outer, + MainloopPipelineReducer& pipeline_reducer, + PipelineStateReducer& smem_pipe_write_reducer, + SharedStorage& storage, + MathWgOrderBarrier& math_wg_order_barrier) { + TiledMmaND tiled_mma_nd; + + Tensor acc_DV = + partition_fragment_C(tiled_mma_nd, take<0, 2>(TileShapeND{})); + clear(acc_DV); + + Tensor acc_DK = + partition_fragment_C(tiled_mma_nd, take<0, 2>(TileShapeND{})); + clear(acc_DK); + + auto [q_offset, q_tiles, qkv_shape, is_fully_block_sparse, has_q_padding] = + get_trip_info(params, blk_coord, problem_size); + int inner_tile_count = size(q_tiles); + + // Dilated cases where dilation doesn't evenly divide input shape can + // sometimes lead to threadblocks with no work to do (similar to varlen). We + // should exit those immediately instead of letting them proceed because + // some of the roles (i.e. load) unroll their loops, and this can cause + // stalls. + if (inner_tile_count == 0) { + return make_tuple(acc_DK, acc_DV); + } + + int thread_idx = int(threadIdx.x) % cutlass::NumThreadsPerWarpGroup; + + PipelineState smem_pipe_release_inner = smem_pipe_read_inner; + + pipeline_outer.consumer_wait(smem_pipe_read_outer); + PipelineStateQ smem_pipe_read_k = smem_pipe_read_outer; + ++smem_pipe_read_outer; + pipeline_outer.consumer_wait(smem_pipe_read_outer); + PipelineStateQ smem_pipe_read_v = smem_pipe_read_outer; + + TiledMmaNM tiled_mma_nm; + Tensor sK = + make_tensor(make_smem_ptr(storage.smem_k.data()), SmemLayoutK{}); + Tensor sQ = + make_tensor(make_smem_ptr(storage.smem_q.data()), SmemLayoutQ{}); + auto thr_mma_nm = tiled_mma_nm.get_thread_slice(thread_idx); + Tensor tSsK = thr_mma_nm.partition_A(sK); + Tensor tSsQ = thr_mma_nm.partition_B(sQ); + Tensor tSrK = thr_mma_nm.make_fragment_A(tSsK); + Tensor tSrQ = thr_mma_nm.make_fragment_B(tSsQ); + + Tensor sV = + make_tensor(make_smem_ptr(storage.smem_v.data()), SmemLayoutV{}); + Tensor sDO = + make_tensor(make_smem_ptr(storage.smem_do.data()), SmemLayoutDO{}); + + Tensor tDPsV = thr_mma_nm.partition_A(sV); + Tensor tDPsDO = thr_mma_nm.partition_B(sDO); + Tensor tDPrV = thr_mma_nm.make_fragment_A(tDPsV); + Tensor tDPrDO = thr_mma_nm.make_fragment_B(tDPsDO); + + auto thr_mma_nd = tiled_mma_nd.get_thread_slice(thread_idx); + + Tensor sDOp = + make_tensor(make_smem_ptr(storage.smem_dop.data()), SmemLayoutDOp{}); + Tensor tDV_sDO = thr_mma_nd.partition_B(sDOp); + Tensor tDVrDO = thr_mma_nd.make_fragment_B(tDV_sDO); + + Tensor sQp = + make_tensor(make_smem_ptr(storage.smem_qp.data()), SmemLayoutQp{}); + Tensor tDK_sQ = thr_mma_nd.partition_B(sQp); + Tensor tDKrQ = thr_mma_nd.make_fragment_B(tDK_sQ); + + int wg_idx = + __shfl_sync(0xffffffff, get<1>(wg_coord) % NumMmaWarpGroups, 0); + + TiledMmaMD tiled_mma_md; + auto thr_mma_md = tiled_mma_md.get_thread_slice(thread_idx); + Tensor sDS = + make_tensor(make_smem_ptr(storage.smem_ds.data()), SmemLayoutDS{}); + Tensor tDQsDS = thr_mma_md.partition_A(sDS); + Tensor tDQrDS_full = thr_mma_md.make_fragment_A(tDQsDS); + Tensor tDQrDS = tDQrDS_full(_, _, _, _); + Tensor sKp = + make_tensor(make_smem_ptr(storage.smem_kp.data()), SmemLayoutKp{}); + Tensor tDQsK = thr_mma_md.partition_B(sKp); + Tensor tDQrK = thr_mma_md.make_fragment_B(tDQsK); + + Tensor sLSE = make_tensor( + make_smem_ptr(storage.smem_lse.data()), + make_shape( + get<0>(TileShapeNM{}), get<1>(TileShapeNM{}), Int{}), + make_stride(_0{}, _1{}, get<1>(TileShapeNM{}))); + Tensor tSsLSE = thr_mma_nm.partition_C(sLSE); + + Tensor sODO = make_tensor( + make_smem_ptr(storage.smem_sumOdO.data()), + make_shape( + get<0>(TileShapeNM{}), get<1>(TileShapeNM{}), Int{}), + make_stride(_0{}, _1{}, get<1>(TileShapeNM{}))); + Tensor tDPsODO = thr_mma_nm.partition_C(sODO); + + Tensor cS = make_identity_tensor(take<0, 2>(TileShapeNM{})); + Tensor tScS = thr_mma_nm.partition_C(cS); + int n_block = get<1>(wg_coord); + tScS.data() = tScS.data() + E<0>{} * n_block * get<0>(TileShapeNM{}); + + // Transpose + Tensor sDSp_full = sDS.compose(make_layout( + make_shape(size<1>(sDS), size<0>(sDS), size<2>(sDS)), + make_stride(size<0>(sDS), _1{}, size<1>(sDS) * size<0>(sDS)))); + Tensor sDSp = sDSp_full(_, _, _); + Tensor tDPsDS = thr_mma_nm.partition_C(sDSp); + + auto thr_mma_nd_ss = TiledMmaND_SS{}.get_thread_slice(thread_idx); + Tensor tDKsDSp = thr_mma_nd_ss.partition_A(sDSp); + + Tensor tDKrDSp = thr_mma_nd_ss.make_fragment_A(tDKsDSp); + + Tensor sDQ = + make_tensor(make_smem_ptr(storage.smem_dq.data()), SmemLayoutDQ{}); + auto tDQsDQ_full = thr_mma_md.partition_C(sDQ); + + auto smem_pipe_read_k_other = smem_pipe_read_k; + smem_pipe_read_k_other.advance(2); + + int k_index = 0; + + bool no_inner_tiles = true; + while (inner_tile_count > 0) { + if (inner_tile_count == 0) + break; + no_inner_tiles = false; + + pipeline_inner.consumer_wait(smem_pipe_read_inner); + PipelineState smem_pipe_read_q = smem_pipe_read_inner; + ++smem_pipe_read_inner; + PipelineState smem_pipe_read_do = smem_pipe_read_inner; + ++smem_pipe_read_inner; + + // GEMM KQ -> S + Tensor acc_S = + partition_fragment_C(tiled_mma_nm, take<0, 2>(TileShapeNM{})); + + warpgroup_fence_operand(acc_S); + warpgroup_arrive(); + gemm_zero_acc( + tiled_mma_nm, + tSrK(_, _, _, smem_pipe_read_k.index()), + tSrQ(_, _, _, smem_pipe_read_q.index()), + acc_S); + warpgroup_commit_batch(); + + pipeline_inner.consumer_wait(smem_pipe_read_do); + + // GEMM VdO -> dP + Tensor acc_DP = + partition_fragment_C(tiled_mma_nm, take<0, 2>(TileShapeNM{})); + + warpgroup_fence_operand(acc_DP); + warpgroup_arrive(); + gemm_zero_acc( + tiled_mma_nm, + tDPrV(_, _, _, smem_pipe_read_v.index()), + tDPrDO(_, _, _, smem_pipe_read_do.index()), + acc_DP); + warpgroup_commit_batch(); + + Tensor reg_LSE = make_fragment_like(acc_S); + for (int i = 0; i < size(reg_LSE); i++) { + reg_LSE(i) = ((ElementAccumulator)std::log2(std::exp(1.0))) * + tSsLSE(_, _, _, smem_pipe_read_q.index())(i); + } + + Tensor reg_ODO = make_fragment_like(acc_S); + if constexpr (decltype(get<0>(TileShape{}) != _128{})::value) { + for (int i = 0; i < size(reg_ODO); i++) { + reg_ODO(i) = tDPsODO(_, _, _, smem_pipe_read_do.index())(i); + } + } + + warpgroup_wait<1>(); + warpgroup_fence_operand(acc_S); + + math_wg_order_barrier.wait(); + // Compute S -> P + + if (not is_fully_block_sparse) { + Mask{}.apply_mask( + acc_S, + tScS, + MultiDimTileShape{}, + params.fna.kv_shape, + qkv_shape, + params.fna.na_params, + q_offset, + q_tiles); + } else if (has_q_padding) { + Mask{}.apply_padded_mask( + acc_S, + tScS, + MultiDimTileShape{}, + qkv_shape, + params.fna.na_params, + q_offset, + q_tiles); + } + + auto acc_P = make_fragment_like(acc_S); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(acc_P); i++) { + acc_P(i) = ::exp2f(params.scale_softmax_log2 * acc_S(i) - reg_LSE(i)); + } + math_wg_order_barrier.arrive(); + + if constexpr (decltype(get<0>(TileShape{}) == _128{})::value) { + for (int i = 0; i < size(reg_ODO); i++) { + reg_ODO(i) = tDPsODO(_, _, _, smem_pipe_read_do.index())(i); + } + } + + warpgroup_wait<0>(); + warpgroup_fence_operand(acc_DP); + + // Compute dP P -> dS + auto acc_DS = make_fragment_like(acc_DP); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(acc_DS); i++) { + // We could move the scale out and into the respective epilogues (or a + // final scaling step) + acc_DS(i) = static_cast( + acc_P(i) * params.scale_softmax * (acc_DP(i) - reg_ODO(i))); + } + + // GEMM PdO -> dV + auto op_P = + make_acc_into_op(acc_P, typename TiledMmaND::LayoutA_TV{}); + warpgroup_fence_operand(acc_DV); + warpgroup_fence_operand(op_P); + warpgroup_arrive(); + cute::gemm( + tiled_mma_nd, + op_P, + tDVrDO(_, _, _, smem_pipe_read_do.index()), + acc_DV); + warpgroup_commit_batch(); + + // Store dS to smem dS' + if (wg_idx == 0) + math_wg_order_barrier.wait(); + + auto recast_bits = [](auto sz, auto t) { + return recast>(t); + }; + auto tDPsDS_v = recast_bits(Int * 2>{}, tDPsDS); + auto acc_DS_v = recast_bits(Int * 2>{}, acc_DS); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(acc_DS_v); i++) { + tDPsDS_v(_, _, _, wg_idx)(i) = acc_DS_v(i); + } + + cutlass::arch::fence_view_async_shared(); + if (wg_idx == 0) + math_wg_order_barrier.arrive(); + + // GEMM dS Q -> dK + if (wg_idx == 1) { + math_wg_order_barrier.wait(); + + // GEMM dS' K -> dQ + Tensor acc_DQ = + partition_fragment_C(tiled_mma_md, take<0, 2>(TileShapeMD{})); + + warpgroup_fence_operand(acc_DQ); + warpgroup_arrive(); + gemm_zero_acc( + tiled_mma_md, + tDQrDS(_, _, _, 0), + tDQrK(_, _, _, smem_pipe_read_k_other.index()), + acc_DQ); + cute::gemm( + tiled_mma_md, + tDQrDS(_, _, _, 1), + tDQrK(_, _, _, smem_pipe_read_k.index()), + acc_DQ); + warpgroup_commit_batch(); + + warpgroup_fence_operand(acc_DK); + warpgroup_arrive(); + cute::gemm( + TiledMmaND_SS{}, + tDKrDSp(_, _, _, wg_idx), + tDKrQ(_, _, _, smem_pipe_read_q.index()), + acc_DK); + warpgroup_commit_batch(); + + warpgroup_wait<1>(); + warpgroup_fence_operand(acc_DK); + + warpgroup_wait<1>(); + warpgroup_fence_operand(acc_DQ); + + math_wg_order_barrier.arrive(); + + pipeline_reducer.producer_acquire(smem_pipe_write_reducer); + auto tDQsDQ = tDQsDQ_full(_, _, _, smem_pipe_write_reducer.index()); + + // Store dQ to smem dQ' + // Invoke TMA reduce on dQ' + using Vec = uint_bit_t * 2>; + auto tDQsDQ_v = recast(tDQsDQ); + auto acc_DQ_v = recast(acc_DQ); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(acc_DQ_v); i++) { + tDQsDQ_v(i) = acc_DQ_v(i); + } + + cutlass::arch::fence_view_async_shared(); + + pipeline_reducer.producer_commit(smem_pipe_write_reducer); + ++smem_pipe_write_reducer; + } else { + warpgroup_fence_operand(acc_DK); + warpgroup_arrive(); + cute::gemm( + TiledMmaND_SS{}, + tDKrDSp(_, _, _, wg_idx), + tDKrQ(_, _, _, smem_pipe_read_q.index()), + acc_DK); + warpgroup_commit_batch(); + + warpgroup_wait<1>(); + warpgroup_fence_operand(acc_DK); + + pipeline_reducer.producer_acquire(smem_pipe_write_reducer); + pipeline_reducer.producer_commit(smem_pipe_write_reducer); + ++smem_pipe_write_reducer; + } + + --inner_tile_count; + + pipeline_inner.consumer_release(smem_pipe_release_inner); + ++smem_pipe_release_inner; + pipeline_inner.consumer_release(smem_pipe_release_inner); + ++smem_pipe_release_inner; + + tScS.data() = tScS.data() + E<1>{} * get<1>(TileShapeNM{}); + k_index += 1; + } + + if (no_inner_tiles) { + // when there's 0 inner tiles, we also need to make sure the order barrier + // stays in sync with the rest of the WG. + + // S -> P + math_wg_order_barrier.wait(); + math_wg_order_barrier.arrive(); + + math_wg_order_barrier.wait(); + math_wg_order_barrier.arrive(); + } + + pipeline_outer.consumer_release(smem_pipe_read_k); + pipeline_outer.consumer_release(smem_pipe_read_outer); + + // When there's 0 inner tiles (i.e. 1 Q tile, > 1 KV tiles, top-left causal + // mask has entire KV tiles to which nothing attends), reduce warp exits + // early and signals no arrivals on pipeline_reducer. If we don't jump over + // producer_tail, acquires will stall indefinitely. + if (not no_inner_tiles) { + pipeline_reducer.producer_tail(smem_pipe_write_reducer); + ++smem_pipe_read_outer; + } + + warpgroup_wait<0>(); + warpgroup_fence_operand(acc_DK); + warpgroup_fence_operand(acc_DV); + + return make_tuple(acc_DK, acc_DV); + } +}; + +} // namespace cutlass::fna::collective diff --git a/natten/csrc/include/natten/cuda/fna_hopper/collective/fna_collective_load.hpp b/natten/csrc/include/natten/cuda/fna_hopper/collective/fna_collective_load.hpp new file mode 100644 index 00000000..3b6fa25c --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna_hopper/collective/fna_collective_load.hpp @@ -0,0 +1,232 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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::fna::collective { + +enum class LoadKind { kQ, kK, kV, kBwdN, kBwdM, kBwdScalar }; + +template < + LoadKind kKind, + class Pipeline, + class Element, + class SmemLayout, + class TMA> +struct CollectiveLoadTma { + using Params = TMA; + using SharedStorage = + cute::array_aligned>; + using PipelineState = typename cutlass::PipelineState; + + Params const& params; + Pipeline& pipeline; + SharedStorage& storage; + + CUTLASS_DEVICE + CollectiveLoadTma( + Params const& params, + Pipeline& pipeline, + SharedStorage& storage) + : params(params), pipeline(pipeline), storage(storage) {} + + template + CUTLASS_DEVICE auto init_g( + ProblemSize const& problem_size, + TileShape const& tile_shape, + BlockCoord const& blk_coord, + int loop_count) { + using X = Underscore; + if constexpr (kKind == LoadKind::kK) { + Tensor mK_full = params.get_tma_tensor(make_shape( + get<3>(problem_size), + get<4>(problem_size), + select<0, 1>(problem_size))); + Tensor gK_full = local_tile( + mK_full, tile_shape, make_coord(_, _, _), Step{}); + Tensor gK = gK_full(_, _, _, _0{}, get<2>(blk_coord)); + return gK; + } else if constexpr (kKind == LoadKind::kQ) { + Tensor mQ_full = params.get_tma_tensor(make_shape( + get<2>(problem_size), + get<4>(problem_size), + select<0, 1>(problem_size))); + Tensor gQ_full = local_tile( + mQ_full, tile_shape, make_coord(_, _, _), Step<_1, X, _1>{}); + Tensor gQ = gQ_full(_, _, _, _0{}, get<2>(blk_coord)); + return make_tensor( + gQ.data() + loop_count * get<0>(blk_coord) * stride<2>(gQ), + gQ.layout()); + } else if constexpr (kKind == LoadKind::kV) { + Tensor mV_full = params.get_tma_tensor(make_shape( + get<4>(problem_size), + get<3>(problem_size), + select<0, 1>(problem_size))); + Tensor gV_full = local_tile( + mV_full, tile_shape, make_coord(_, _, _), Step{}); + Tensor gV = gV_full(_, _, _0{}, _, get<2>(blk_coord)); + return gV; + } else if constexpr (kKind == LoadKind::kBwdN) { + Tensor m_full = params.get_tma_tensor(make_shape( + get<3>(problem_size), + get<4>(problem_size), + select<0, 1>(problem_size))); + Tensor g_full = local_tile( + m_full, tile_shape, make_coord(_, _, _), Step<_1, X, _1>{}); + Tensor g = g_full(_, _, _, _0{}, get<2>(blk_coord)); + return make_tensor( + g.data() + loop_count * get<1>(blk_coord) * stride<2>(g), g.layout()); + } else if constexpr (kKind == LoadKind::kBwdM) { + Tensor m_full = params.get_tma_tensor(make_shape( + get<2>(problem_size), + get<4>(problem_size), + select<0, 1>(problem_size))); + Tensor g_full = local_tile( + m_full, tile_shape, make_coord(_, _, _), Step{}); + Tensor g = g_full(_, _, _, _0{}, get<2>(blk_coord)); + return g; + } else if constexpr (kKind == LoadKind::kBwdScalar) { + Tensor m_full = params.get_tma_tensor(select<2, 0, 1>(problem_size)); + Tensor g_full = + local_tile(m_full, tile_shape, make_coord(_, _, _), Step{}); + Tensor g = g_full(_, _, get<2, 0>(blk_coord), get<2, 1>(blk_coord)); + return g; + } + } + + template < + class ClusterRank, + class ProblemSize, + class TileShape, + class BlockCoord> + CUTLASS_DEVICE auto init_state( + ClusterRank const& block_rank_in_cluster, + ProblemSize const& problem_size, + TileShape const& tile_shape, + BlockCoord const& block_coord, + int loop_count) { + Tensor g = init_g(problem_size, tile_shape, block_coord, loop_count); + Tensor s = make_tensor(make_smem_ptr(storage.data()), SmemLayout{}); + + auto block_tma = params.get_slice(block_rank_in_cluster); + Tensor ts = block_tma.partition_D(s); + Tensor tg = block_tma.partition_S(g); + + return make_tuple(tg, ts); + } + + template < + bool kAdvanceIterator = true, + bool kAdvancePipe = true, + bool kAcquireBarrier = true, + class TileIterator, + class State, + class IterToTileMap> + CUTLASS_DEVICE void step_with_iter_tile_map( + TileIterator& tile_iter, + State const& state, + PipelineState& smem_pipe_write, + int lane_predicate, + int& tile_count, + uint16_t mcast_mask, + IterToTileMap const& iter_to_tile_map) { + if ((lane_predicate == 1) && (tile_count > 0)) { + if constexpr (kAcquireBarrier) + pipeline.producer_acquire(smem_pipe_write); + using BarrierType = typename Pipeline::ProducerBarrierType; + BarrierType* tma_barrier = pipeline.producer_get_barrier(smem_pipe_write); + + if constexpr (kKind == LoadKind::kBwdScalar) { + copy( + params.with(*tma_barrier, mcast_mask), + get<0>(state)(_, _, iter_to_tile_map(*tile_iter)), + get<1>(state)(_, _, smem_pipe_write.index())); + } else { + copy( + params.with(*tma_barrier, mcast_mask), + get<0>(state)(_, _, _, iter_to_tile_map(*tile_iter)), + get<1>(state)(_, _, _, smem_pipe_write.index())); + } + + if constexpr (kAdvancePipe) + ++smem_pipe_write; + if constexpr (kAdvanceIterator) + ++tile_iter; + } + + --tile_count; + } + + template < + bool kAdvanceIterator = true, + bool kAdvancePipe = true, + bool kAcquireBarrier = true, + class TileIterator, + class State> + CUTLASS_DEVICE void step( + TileIterator& tile_iter, + State const& state, + PipelineState& smem_pipe_write, + int lane_predicate, + int& tile_count, + uint16_t mcast_mask = 0) { + if ((lane_predicate == 1) && (tile_count > 0)) { + if constexpr (kAcquireBarrier) + pipeline.producer_acquire(smem_pipe_write); + using BarrierType = typename Pipeline::ProducerBarrierType; + BarrierType* tma_barrier = pipeline.producer_get_barrier(smem_pipe_write); + + if constexpr (kKind == LoadKind::kBwdScalar) { + copy( + params.with(*tma_barrier, mcast_mask), + get<0>(state)(_, _, *tile_iter), + get<1>(state)(_, _, smem_pipe_write.index())); + } else { + copy( + params.with(*tma_barrier, mcast_mask), + get<0>(state)(_, _, _, *tile_iter), + get<1>(state)(_, _, _, smem_pipe_write.index())); + } + + if constexpr (kAdvancePipe) + ++smem_pipe_write; + if constexpr (kAdvanceIterator) + ++tile_iter; + } + --tile_count; + } +}; + +} // namespace cutlass::fna::collective diff --git a/natten/csrc/include/natten/cuda/fna_hopper/collective/fna_collective_softmax.hpp b/natten/csrc/include/natten/cuda/fna_hopper/collective/fna_collective_softmax.hpp new file mode 100644 index 00000000..b631003a --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna_hopper/collective/fna_collective_softmax.hpp @@ -0,0 +1,479 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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" + +#include "natten/cuda/fna_hopper/collective/fna_common.hpp" + +namespace cutlass::fna::collective { + +template +struct CollectiveSoftmax { + Params const& params; + CUTLASS_DEVICE CollectiveSoftmax(Params const& params) : params(params) {} + + using SumType = float; + using MaxType = ElementAccumulator; + + template + CUTLASS_DEVICE auto init( + AccPV const& acc_pv, + TiledMmaPV const& tiled_mma_pv) { + Tensor s_max = make_fragment_like( + size<0>(layout_acc_mn(tiled_mma_pv, acc_pv.layout()))); + Tensor a_sum = make_fragment_like(s_max); + return make_tuple(s_max, a_sum); + } + + CUTLASS_DEVICE float overload_exp2(float f) { + return ::exp2f(f); + } + + CUTLASS_DEVICE cutlass::half_t overload_exp2(cutlass::half_t f) { + auto a = f.raw(); + decltype(a) d; + asm("ex2.approx.f16 %0, %1;" : "=h"(d) : "h"(a)); + return cutlass::half_t::bitcast(d); + } + + CUTLASS_DEVICE float overload_max(float a, float b) { + return ::max(a, b); + } + + CUTLASS_DEVICE cutlass::half_t overload_max( + cutlass::half_t a, + cutlass::half_t b) { + return cutlass::half_t{__hmax_nan(a.to_half(), b.to_half())}; + } + + CUTLASS_DEVICE half overload_to_native(cutlass::half_t f) { + return f.to_half(); + } + + CUTLASS_DEVICE float overload_to_native(float f) { + return f; + } + + template < + class AccQK, + class TiledMmaQK, + class CountQK, + class State, + class ProblemShape, + class NADim, + class NAParams, + class MultiDimTileShape> + CUTLASS_DEVICE auto step( + AccQK& acc_qk, + TiledMmaQK const& tiled_mma_qk, + CountQK const& count_qk, + State& state, + ProblemShape const& problem_shape, + bool is_fully_block_sparse, + bool has_kv_padding, + NADim const& q_shape, + NADim const& qkv_shape, + NAParams const& na_params, + NADim const& blk_kv_offset, + NADim const& kv_diff_tiles, + MultiDimTileShape const& multi_dim_tile_shape) { + // Fusion{}.before_softmax(acc_qk, count_qk, problem_shape); + Tensor acc_qk_mn = make_tensor( + acc_qk.data(), layout_acc_mn(tiled_mma_qk, acc_qk.layout())); + + if (not is_fully_block_sparse) { + Fusion{}.apply_mask( + acc_qk_mn, + count_qk, + multi_dim_tile_shape, + q_shape, + qkv_shape, + na_params, + blk_kv_offset, + kv_diff_tiles); + } else if (has_kv_padding) { + Fusion{}.apply_padded_mask( + acc_qk_mn, + count_qk, + multi_dim_tile_shape, + qkv_shape, + na_params, + blk_kv_offset, + kv_diff_tiles); + } + + auto reduction_target_qk = reduction_target_n(tiled_mma_qk); + constexpr int red_rank = decltype(rank(reduction_target_qk))::value; + + auto& s_max = get<0>(state); + auto& a_sum = get<1>(state); + + // Linear reduction is faster for the first iteration + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<0>(acc_qk_mn); i++) { + s_max(i) = acc_qk_mn(i, 0); + } + CUTLASS_PRAGMA_UNROLL + for (int j = 1; j < size<1>(acc_qk_mn); j++) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<0>(acc_qk_mn); i++) { + s_max(i) = overload_max(s_max(i), acc_qk_mn(i, j)); + } + } + + for_each(make_seq{}, [&](auto r) { + CUTLASS_PRAGMA_UNROLL + for (int j = 1; j < shape(reduction_target_qk); j *= 2) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<0>(acc_qk_mn); i++) { + s_max(i) = overload_max( + s_max(i), + MaxType{__shfl_xor_sync( + uint32_t(-1), + overload_to_native(s_max(i)), + stride(reduction_target_qk) * j)}); + } + } + }); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<0>(acc_qk_mn); i++) { + MaxType local_max = s_max(i) == static_cast(-INFINITY) + ? static_cast(0) + : s_max(i); + MaxType scale = static_cast(params.scale_softmax_log2); + MaxType scale_max = scale * local_max; + + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<1>(acc_qk_mn); j++) { + acc_qk_mn(i, j) = overload_exp2(scale * acc_qk_mn(i, j) - scale_max); + } + } + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<0>(acc_qk_mn); i++) { + a_sum(i) = SumType{reduce(acc_qk_mn(i, _), cute::plus{})}; + } + } + + template < + bool kUseFusion = true, + class AccQK, + class TiledMmaQK, + class CountQK, + class State, + class AccPV, + class TiledMmaPV, + class ProblemShape, + class NADim, + class NAParams, + class MultiDimTileShape> + CUTLASS_DEVICE auto step_interleave_begin( + AccQK& acc_qk, + TiledMmaQK const& tiled_mma_qk, + CountQK const& count_qk, + State& state, + AccPV& acc_pv, + TiledMmaPV const& tiled_mma_pv, + ProblemShape const& problem_shape, + bool is_fully_block_sparse, + bool has_kv_padding, + NADim const& q_shape, + NADim const& qkv_shape, + NAParams const& na_params, + NADim const& blk_kv_offset, + NADim const& kv_diff_tiles, + MultiDimTileShape const& multi_dim_tile_shape) { + // if constexpr (kUseFusion) { + // Fusion{}.before_softmax(acc_qk, count_qk, problem_shape); + // } + + Tensor acc_qk_mn = make_tensor( + acc_qk.data(), layout_acc_mn(tiled_mma_qk, acc_qk.layout())); + Tensor acc_pv_mn = make_tensor( + acc_pv.data(), layout_acc_mn(tiled_mma_pv, acc_pv.layout())); + + if constexpr (kUseFusion) { + if (not is_fully_block_sparse) { + Fusion{}.apply_mask( + acc_qk_mn, + count_qk, + multi_dim_tile_shape, + q_shape, + qkv_shape, + na_params, + blk_kv_offset, + kv_diff_tiles); + } else if (has_kv_padding) { + Fusion{}.apply_padded_mask( + acc_qk_mn, + count_qk, + multi_dim_tile_shape, + qkv_shape, + na_params, + blk_kv_offset, + kv_diff_tiles); + } + } + + static_assert(size<0>(acc_qk_mn) == size<0>(acc_pv_mn)); + auto reduction_target_qk = reduction_target_n(tiled_mma_qk); + constexpr int red_rank = decltype(rank(reduction_target_qk))::value; + + auto& s_max = get<0>(state); + auto& a_sum = get<1>(state); + + Tensor s_max_prev = make_fragment_like(s_max); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<0>(acc_qk_mn); i++) { + s_max_prev(i) = s_max(i); + } + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<0>(acc_qk_mn); i++) { + // Linear reduction is faster here, as well + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<1>(acc_qk_mn); j++) { + s_max(i) = overload_max(s_max(i), acc_qk_mn(i, j)); + } + } + // reduce max + for_each(make_seq{}, [&](auto r) { + CUTLASS_PRAGMA_UNROLL + for (int j = 1; j < shape(reduction_target_qk); j *= 2) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<0>(acc_qk_mn); i++) { + s_max(i) = overload_max( + s_max(i), + __shfl_xor_sync( + uint32_t(-1), s_max(i), stride(reduction_target_qk) * j)); + } + } + }); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<0>(acc_pv_mn); i++) { + float s_max_cur = s_max(i) == -INFINITY ? 0.0f : s_max(i); + float scale = + ::exp2f((s_max_prev(i) - s_max_cur) * params.scale_softmax_log2); + a_sum(i) *= scale; + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<1>(acc_pv_mn); j++) { + acc_pv_mn(i, j) *= scale; + } + } + } + + template + CUTLASS_DEVICE auto step_interleave_step(AccQK_MN& acc_qk_mn, State& state) { + auto& s_max = get<0>(state); + auto& a_sum = get<1>(state); + + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<0>(acc_qk_mn); j++) { + float local_max = s_max(j) == -INFINITY ? 0.f : s_max(j); + float scale_max = params.scale_softmax_log2 * local_max; + + CUTLASS_PRAGMA_UNROLL + for (int k = 0; k < size<1>(acc_qk_mn); k++) { + acc_qk_mn(j, k) = + ::exp2f(params.scale_softmax_log2 * acc_qk_mn(j, k) - scale_max); + a_sum(j) += acc_qk_mn(j, k); + } + } + } + + template < + bool kUseFusion = true, + class AccQK, + class TiledMmaQK, + class CountQK, + class State, + class AccPV, + class TiledMmaPV, + class ProblemShape, + class NADim, + class NAParams, + class MultiDimTileShape> + CUTLASS_DEVICE auto step( + AccQK& acc_qk, + TiledMmaQK const& tiled_mma_qk, + CountQK const& count_qk, + State& state, + AccPV& acc_pv, + TiledMmaPV const& tiled_mma_pv, + ProblemShape const& problem_shape, + bool is_fully_block_sparse, + bool has_kv_padding, + NADim const& q_shape, + NADim const& qkv_shape, + NAParams const& na_params, + NADim const& blk_kv_offset, + NADim const& kv_diff_tiles, + MultiDimTileShape const& multi_dim_tile_shape) { + // if constexpr (kUseFusion) { + // Fusion{}.before_softmax(acc_qk, count_qk, problem_shape); + // } + + Tensor acc_qk_mn = make_tensor( + acc_qk.data(), layout_acc_mn(tiled_mma_qk, acc_qk.layout())); + Tensor acc_pv_mn = make_tensor( + acc_pv.data(), layout_acc_mn(tiled_mma_pv, acc_pv.layout())); + + if constexpr (kUseFusion) { + if (not is_fully_block_sparse) { + Fusion{}.apply_mask( + acc_qk_mn, + count_qk, + multi_dim_tile_shape, + q_shape, + qkv_shape, + na_params, + blk_kv_offset, + kv_diff_tiles); + } else if (has_kv_padding) { + Fusion{}.apply_padded_mask( + acc_qk_mn, + count_qk, + multi_dim_tile_shape, + qkv_shape, + na_params, + blk_kv_offset, + kv_diff_tiles); + } + } + + static_assert(size<0>(acc_qk_mn) == size<0>(acc_pv_mn)); + auto reduction_target_qk = reduction_target_n(tiled_mma_qk); + constexpr int red_rank = decltype(rank(reduction_target_qk))::value; + + auto& s_max = get<0>(state); + auto& a_sum = get<1>(state); + + Tensor s_max_prev = make_fragment_like(s_max); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<0>(acc_qk_mn); i++) { + s_max_prev(i) = s_max(i); + + // Linear reduction is faster here, as well + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<1>(acc_qk_mn); j++) { + s_max(i) = overload_max(s_max(i), acc_qk_mn(i, j)); + } + // reduce max + for_each(make_seq{}, [&](auto r) { + CUTLASS_PRAGMA_UNROLL + for (int j = 1; j < shape(reduction_target_qk); j *= 2) { + s_max(i) = overload_max( + s_max(i), + MaxType{__shfl_xor_sync( + uint32_t(-1), + overload_to_native(s_max(i)), + stride(reduction_target_qk) * j)}); + } + }); + + MaxType local_max = s_max(i) == static_cast(-INFINITY) + ? static_cast(0) + : s_max(i); + MaxType scale = static_cast(params.scale_softmax_log2); + MaxType scale_max = scale * local_max; + + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<1>(acc_qk_mn); j++) { + acc_qk_mn(i, j) = overload_exp2(scale * acc_qk_mn(i, j) - scale_max); + } + + MaxType s_max_cur = s_max(i) == static_cast(-INFINITY) + ? static_cast(0) + : s_max(i); + SumType scale_pv = overload_exp2((s_max_prev(i) - s_max_cur) * scale); + a_sum(i) *= scale_pv; + + using ElementPV = typename AccPV::value_type; + ElementPV scale_pv_ele = static_cast(scale_pv); + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<1>(acc_pv_mn); j++) { + acc_pv_mn(i, j) *= scale_pv_ele; + } + a_sum(i) += SumType{reduce(acc_qk_mn(i, _), cute::plus{})}; + } + } + + template + CUTLASS_DEVICE auto tail( + State& state, + AccPV& acc_pv, + TiledMmaPV const& tiled_mma_pv) { + auto& s_max = get<0>(state); + auto& a_sum = get<1>(state); + + Tensor acc_pv_mn = make_tensor( + acc_pv.data(), layout_acc_mn(tiled_mma_pv, acc_pv.layout())); + + auto reduction_target = reduction_target_n(tiled_mma_pv); + constexpr int red_rank = decltype(rank(reduction_target))::value; + for_each(make_seq{}, [&](auto r) { + CUTLASS_PRAGMA_UNROLL + for (int j = 1; j < shape(reduction_target); j *= 2) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<0>(acc_pv_mn); i++) { + a_sum(i) = + a_sum(i) + + __shfl_xor_sync( + uint32_t(-1), a_sum(i), stride(reduction_target) * j); + } + } + }); + + Tensor acc_mn = make_tensor( + acc_pv.data(), layout_acc_mn(tiled_mma_pv, acc_pv.layout())); + + Tensor lse = make_fragment_like(a_sum); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<0>(acc_mn); i++) { + float sum = a_sum(i); + float inv_sum = (sum == 0.f || sum != sum) ? 1.f : __frcp_rn(sum); + lse(i) = (sum == 0.f || sum != sum) + ? INFINITY + : s_max(i) * params.scale_softmax + __logf(sum); + float scale = params.rp_dropout * inv_sum; + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<1>(acc_mn); j++) { + acc_mn(i, j) *= scale; + } + } + + return lse; + } +}; + +} // namespace cutlass::fna::collective diff --git a/natten/csrc/include/natten/cuda/fna_hopper/collective/fna_collective_tma.hpp b/natten/csrc/include/natten/cuda/fna_hopper/collective/fna_collective_tma.hpp new file mode 100644 index 00000000..55d6071a --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna_hopper/collective/fna_collective_tma.hpp @@ -0,0 +1,828 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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/gemm/collective/collective_builder.hpp" + +#include "natten/cuda/fna_hopper/collective/fna_collective_load.hpp" +#include "natten/cuda/fna_hopper/collective/fna_collective_softmax.hpp" +#include "natten/cuda/fna_hopper/collective/fna_common.hpp" + +// This is identical to FMHA; reuse +#include "natten/cuda/fmha_hopper/kernel/fmha_options.hpp" + +namespace cutlass::fna::collective { + +using namespace cute; +using cutlass::fmha::kernel::find_option_t; +using cutlass::fmha::kernel::Tag; + +template < + class NADim, + class QTileShape, + class KVTileShape, + typename Element_, + typename ElementAccumulator_, + typename TileShape_, // BlockQO, BlockKV, BlockHead + class Fusion, + class... Options> +struct FnaMainloopTmaSm90 { + using Element = Element_; + using ElementAccumulator = ElementAccumulator_; + using TileShape = TileShape_; + + using MultiDimTileShape = cute::tuple; + + static_assert( + size(QTileShape{}) == get<0>(TileShape{}), + "QTileShape doesn't match the size of Q tile in the FMHA kernel."); + static_assert( + size(KVTileShape{}) == get<1>(TileShape{}), + "QTileShape doesn't match the size of Q tile in the FMHA kernel."); + + // Options + using kClusterM = find_option_t, Options...>; + static constexpr int StageCount = + find_option_t, Options...>::value; + static constexpr int StageCountQ = + find_option_t, Options...>::value; + + using StagesQ = cutlass::gemm::collective::StageCount; + using Stages = cutlass::gemm::collective::StageCount; + using ClusterShape = Shape; + + // 16B alignment lets us use TMA + static constexpr int Alignment = 16 / sizeof(Element); + + using TileShapeQK = TileShape; + using TileShapePV = decltype(select<0, 2, 1>(TileShapeQK{})); + + using LayoutQKV = cute::tuple>; + using LayoutQ = LayoutQKV; + using LayoutK = LayoutQKV; + using LayoutV = LayoutQKV; + + using CollectiveMmaQK = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm90, + cutlass::arch::OpClassTensorOp, + Element, + LayoutQ, + Alignment, + Element, + LayoutK, + Alignment, + ElementAccumulator, + TileShapeQK, + ClusterShape, + Stages, + cutlass::gemm::KernelTmaWarpSpecialized>::CollectiveOp; + + using CollectiveMmaPV = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm90, + cutlass::arch::OpClassTensorOp, + // the stride for A does not matter since we do not load from smem at all + Element, + LayoutK, + Alignment, + Element, + decltype(select<1, 0, 2>(LayoutV{})), + Alignment, + ElementAccumulator, + TileShapePV, + ClusterShape, + Stages, + cutlass::gemm::KernelTmaWarpSpecialized>::CollectiveOp; + + using TiledMmaQK = typename CollectiveMmaQK::TiledMma; + using TiledMmaPV = + decltype(convert_to_gmma_rs(typename CollectiveMmaPV::TiledMma{})); + + using SmemLayoutQ = decltype(unstageSmemLayout( + typename CollectiveMmaQK::SmemLayoutA{}, + Int{})); + using SmemLayoutK = typename CollectiveMmaQK::SmemLayoutB; + using SmemLayoutV = typename CollectiveMmaPV::SmemLayoutB; + + using MainloopPipeline = cutlass::PipelineTmaAsync; + using MainloopPipelineQ = cutlass::PipelineTmaAsync; + + using PipelineState = + typename cutlass::PipelineState; + using PipelineStateQ = + typename cutlass::PipelineState; + + using TileShapeOut = TileShapePV; + using TiledMmaOut = TiledMmaPV; + using ElementOut = ElementAccumulator; + + struct SharedStorage { + cute::array_aligned> smem_q; + union { + cute::array_aligned> smem_k; + cute::array_aligned> smem_v; + }; + }; + + struct Arguments { + const Element* ptr_Q; + LayoutQ dQ; + const Element* ptr_K; + LayoutK dK; + const Element* ptr_V; + LayoutV dV; + + // if zero, defaults to 1/sqrt(D) + float scale_softmax = 0.0f; + + // FNA params + NADim q_shape; + NADim kv_shape; + NADim qkv_shape; + NADim window_size; + NADim stride; + NADim dilation; + }; + + 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; + TMA_K tma_load_k; + TMA_V tma_load_v; + + float scale_softmax; + float scale_softmax_log2; + float rp_dropout; + + // FNA params + NADim qkv_shape; + NADim q_shape; + NADim kv_shape; + cute::tuple + na_params; // win, win_left, win_right, stride + bool is_fully_block_sparse; + bool has_kv_padding; + NADim dilation; + bool requires_qkv_fixup; + bool is_dilated; + int num_dilation_groups; + }; + + using LoadQ = cutlass::fna::collective::CollectiveLoadTma< + cutlass::fna::collective::LoadKind::kQ, + MainloopPipelineQ, + Element, + SmemLayoutQ, + TMA_Q>; + + using LoadK = cutlass::fna::collective::CollectiveLoadTma< + cutlass::fna::collective::LoadKind::kK, + MainloopPipeline, + Element, + SmemLayoutK, + TMA_K>; + + using LoadV = cutlass::fna::collective::CollectiveLoadTma< + cutlass::fna::collective::LoadKind::kV, + MainloopPipeline, + Element, + SmemLayoutV, + TMA_V>; + + static_assert( + size(typename CollectiveMmaQK::TiledMma{}) == + size(typename CollectiveMmaPV::TiledMma{})); + + static const int MaxThreadsPerBlock = + size(typename CollectiveMmaQK::TiledMma{}); + + template + static bool can_implement( + ProblemShape const& problem_size, + Arguments const& args) { + return (get<4>(problem_size) <= get<2>(TileShape{})) && + ((get<4>(problem_size) % Alignment) == 0) && + ((get<2>(problem_size) % Alignment) == 0) && + // FNA + evenly_divides(args.q_shape, QTileShape{}) && + evenly_divides(args.kv_shape, KVTileShape{}) && + evenly_divides( + size<0>(problem_size), + size(args.dilation)) && // dilation groups are + // folded into batch + tuple_leq(args.window_size, args.qkv_shape) && + // TODO: check window size * dilation <= qkv shape + tuple_leq(args.stride, args.window_size); + } + + template + static Params to_underlying_arguments( + ProblemShape const& problem_size, + Arguments const& args, + void* workspace) { + auto problem_shape_qk = make_shape( + get<2>(problem_size), + get<3>(problem_size), + get<4>(problem_size), + make_shape(get<0>(problem_size), get<1>(problem_size))); + auto params_qk = CollectiveMmaQK::to_underlying_arguments( + problem_shape_qk, + typename CollectiveMmaQK::Arguments{ + args.ptr_Q, + args.dQ, + args.ptr_K, + args.dK, + }, + /*workspace=*/nullptr); + + auto problem_shape_pv = select<0, 2, 1, 3>(problem_shape_qk); + auto params_pv = CollectiveMmaPV::to_underlying_arguments( + problem_shape_pv, + typename CollectiveMmaPV::Arguments{ + args.ptr_K, + args.dK, // never used, dummy + args.ptr_V, + select<1, 0, 2>(args.dV), + }, + /*workspace=*/nullptr); + + float scale_softmax = args.scale_softmax; + if (scale_softmax == 0.0f) { + scale_softmax = 1.0f / (float)std::sqrt(get<4>(problem_size)); + } + float log2_e = static_cast(std::log2(std::exp(1.0))); + + auto window_left = get_window_left(args.window_size); + auto window_right = get_window_right(args.window_size); + + bool requires_qkv_fixup = not evenly_divides(args.qkv_shape, args.dilation); + + return Params{ + params_qk.tma_load_a, + params_qk.tma_load_b, + params_pv.tma_load_b, + scale_softmax, + log2_e * scale_softmax, + 1.0f, + args.qkv_shape, + args.q_shape, + args.kv_shape, + make_tuple(args.window_size, window_left, window_right, args.stride), + fully_block_sparse( + args.qkv_shape, + args.window_size, + args.stride, + QTileShape{}, + KVTileShape{}), + /* has_kv_padding */ not evenly_divides(args.qkv_shape, KVTileShape{}), + args.dilation, + requires_qkv_fixup, + is_dilated(args.dilation), + size(args.dilation)}; + } + + 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 auto compute( + int block_rank_in_cluster, + BlkCoord const& blk_coord, + Params const& params, + ProblemShape const& problem_size, + MainloopPipeline& pipeline, + PipelineState& smem_pipe_read, + PipelineState& smem_pipe_write, + MainloopPipelineQ& pipeline_q, + PipelineStateQ& smem_pipe_read_q, + PipelineStateQ& smem_pipe_write_q, + SharedStorage& storage) { + int warp_idx = cutlass::canonical_warp_idx_sync(); + int thread_idx = threadIdx.x; + + PipelineState smem_pipe_release = smem_pipe_read; + [[maybe_unused]] PipelineStateQ smem_pipe_release_q = smem_pipe_read_q; + + // int fusion_tile_count = + // Fusion{}.get_trip_count(blk_coord, TileShape{}, problem_size); + + auto qkv_shape = params.qkv_shape; + bool is_fully_block_sparse = params.is_fully_block_sparse; + bool has_kv_padding = params.has_kv_padding; + if (params.requires_qkv_fixup) { + qkv_shape = Fusion{}.correct_qkv_shape( + params.qkv_shape, + blk_coord, + params.dilation, + params.num_dilation_groups); + is_fully_block_sparse = fully_block_sparse( + qkv_shape, + get<0>(params.na_params), + get<3>(params.na_params), + QTileShape{}, + KVTileShape{}); + has_kv_padding = not evenly_divides(qkv_shape, KVTileShape{}); + } else if (params.is_dilated) { + qkv_shape = ceil_div(params.qkv_shape, params.dilation); + is_fully_block_sparse = fully_block_sparse( + qkv_shape, + get<0>(params.na_params), + get<3>(params.na_params), + QTileShape{}, + KVTileShape{}); + has_kv_padding = not evenly_divides(qkv_shape, KVTileShape{}); + } + + auto [kv_start, num_tiles] = Fusion{}.get_trip_count( + blk_coord, + MultiDimTileShape{}, + params.q_shape, + qkv_shape, + params.na_params); + + int fusion_tile_count = size(num_tiles); + + auto kv_start_tile = ceil_div(kv_start, KVTileShape{}); + + auto kv_tiled = ceil_div(params.kv_shape, KVTileShape{}); + auto ctr = make_identity_tensor(num_tiles); + auto ctr_offset = domain_offset(kv_start_tile, ctr); + + auto kv_tiled_layout = make_layout(kv_tiled); + + auto iter_to_tile_map = [&ctr_offset, &kv_tiled_layout](int iter) { + return crd2idx(ctr_offset(iter), kv_tiled_layout); + }; + + LoadQ load_q{params.tma_load_q, pipeline_q, storage.smem_q}; + auto load_state_q = + load_q.init_state(_0{}, problem_size, TileShapeQK{}, blk_coord, 1); + + LoadK load_k{params.tma_load_k, pipeline, storage.smem_k}; + auto load_state_k = load_k.init_state( + block_rank_in_cluster, + problem_size, + TileShapeQK{}, + blk_coord, + fusion_tile_count); + + LoadV load_v{params.tma_load_v, pipeline, storage.smem_v}; + auto load_state_v = load_v.init_state( + block_rank_in_cluster, + problem_size, + TileShapePV{}, + blk_coord, + fusion_tile_count); + + // Set predicate for the lowest lane_id in the warp + int lane_predicate = cute::elect_one_sync(); + + // Issue TmaLoads (Prologue fetches) + if (warp_idx == 0) { + auto q_tile_iter = cute::make_coord_iterator(1); + int q_tile_count = 1; + load_q.step( + q_tile_iter, + load_state_q, + smem_pipe_write_q, + lane_predicate, + q_tile_count); + } + + // Loop over K elems + auto k_tile_iter = cute::make_coord_iterator(fusion_tile_count); + + int k_tile_count_tma = 2 * fusion_tile_count; + + uint16_t mcast_mask_b = 0; + + if (warp_idx == 0 && lane_predicate == 1) { + if constexpr (cute::is_same_v< + typename CollectiveMmaQK::GmemTiledCopyB, + SM90_TMA_LOAD_MULTICAST>) { + auto block_layout = Layout{}; // (m,n) -> block_id + for (int m = 0; m < size<0>(block_layout); ++m) { + mcast_mask_b |= (uint16_t(1) << block_layout(m, _0{}, Int<0>{})); + } + } + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < StageCount; i++) { + if (i % 2 == 0) { + load_k.template step_with_iter_tile_map( + k_tile_iter, + load_state_k, + smem_pipe_write, + lane_predicate, + k_tile_count_tma, + mcast_mask_b, + iter_to_tile_map); + } else { + load_v.template step_with_iter_tile_map( + k_tile_iter, + load_state_k, + smem_pipe_write, + lane_predicate, + k_tile_count_tma, + mcast_mask_b, + iter_to_tile_map); + } + } + } + + TiledMmaQK tiled_mma_qk; + auto thr_mma_qk = tiled_mma_qk.get_thread_slice(thread_idx); + + // Mainloop setup QK + 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 tSsQ = thr_mma_qk.partition_A(sQ); // (MMA,MMA_M,MMA_K,PIPE) + Tensor tSsK = thr_mma_qk.partition_B(sK); // (MMA,MMA_N,MMA_K,PIPE) + Tensor tSrQ = thr_mma_qk.make_fragment_A(tSsQ); // (MMA,MMA_N,MMA_K,PIPE) + Tensor tSrK = thr_mma_qk.make_fragment_B(tSsK); // (MMA,MMA_M,MMA_N,PIPE) + + // Prepare: MMA PV + TiledMmaPV tiled_mma_pv; + auto thr_mma_pv = tiled_mma_pv.get_thread_slice(thread_idx); + + // Mainloop setup PV + Tensor sV = + make_tensor(make_smem_ptr(storage.smem_v.data()), SmemLayoutV{}); + + Tensor tOsV = thr_mma_pv.partition_B(sV); // (MMA,MMA_N,MMA_K,PIPE) + Tensor tOrV = thr_mma_pv.make_fragment_B(tOsV); // (MMA,MMA_M,MMA_N,PIPE) + + int k_tile_count = 0; + + pipeline_q.consumer_wait(smem_pipe_read_q); + + // mapping into QK accumulator + Tensor cP = make_identity_tensor(take<0, 2>(TileShapeQK{})); + Tensor tPcP = thr_mma_qk.partition_C(cP); + int m_block = get<0>(blk_coord); + tPcP.data() = tPcP.data() + E<0>{} * m_block * get<0>(TileShapeQK{}); + + // Allocate PV acc + Tensor acc_pv = + partition_fragment_C(tiled_mma_pv, take<0, 2>(TileShapePV{})); + + cutlass::fna::collective:: + CollectiveSoftmax + softmax{params}; + auto softmax_state = softmax.init(acc_pv, tiled_mma_pv); + + if (true) { + --k_tile_count; + // Allocate QK acc + Tensor acc_qk = + partition_fragment_C(tiled_mma_qk, take<0, 2>(TileShapeQK{})); + + pipeline.consumer_wait(smem_pipe_read); + + // MMA QK + warpgroup_fence_operand(acc_qk); + warpgroup_arrive(); + + gemm_zero_acc( + tiled_mma_qk, + tSrQ(_, _, _, _0{}), + tSrK(_, _, _, smem_pipe_read.index()), + acc_qk); + warpgroup_commit_batch(); + + ++smem_pipe_read; + + // Wait for the pipeline MMAs to drain + warpgroup_wait<0>(); + warpgroup_fence_operand(acc_qk); + + softmax.step( + acc_qk, + tiled_mma_qk, + tPcP, + softmax_state, + problem_size, + is_fully_block_sparse, + has_kv_padding, + params.q_shape, + qkv_shape, + params.na_params, + kv_start, + num_tiles, + MultiDimTileShape{}); + + Tensor acc_qk_fixed = + make_fragment_like(convert_c_layout_to_a_layout( + acc_qk.layout(), + shape<1>(typename decltype(tiled_mma_pv)::LayoutA_TV{}))); + + Tensor acc_qk_input = make_tensor(acc_qk_fixed.data(), acc_qk.layout()); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(acc_qk); i++) { + acc_qk_input(i) = static_cast(acc_qk(i)); + } + + pipeline.consumer_wait(smem_pipe_read); + + // MMA PV + warpgroup_fence_operand(acc_pv); + warpgroup_fence_operand(acc_qk_fixed); + warpgroup_arrive(); + + gemm_zero_acc( + tiled_mma_pv, + acc_qk_fixed, + tOrV(_, _, _, smem_pipe_read.index()), + acc_pv); + warpgroup_commit_batch(); + + // + // Advance the pipe + // + + // Advance consumer pipeline + ++smem_pipe_read; + + pipeline.consumer_release(smem_pipe_release); + ++smem_pipe_release; + + tPcP.data() = tPcP.data() + E<1>{} * get<1>(TileShapeQK{}); + } + + CUTLASS_PRAGMA_NO_UNROLL + for (; k_tile_count > 0; --k_tile_count) { + // Allocate QK acc + Tensor acc_qk = + partition_fragment_C(tiled_mma_qk, take<0, 2>(TileShapeQK{})); + + pipeline.consumer_wait(smem_pipe_read); + + // MMA QK + warpgroup_fence_operand(acc_qk); + warpgroup_arrive(); + + gemm_zero_acc( + tiled_mma_qk, + tSrQ(_, _, _, _0{}), + tSrK(_, _, _, smem_pipe_read.index()), + acc_qk); + warpgroup_commit_batch(); + + ++smem_pipe_read; + + if (warp_idx == 0) { + load_k.template step_with_iter_tile_map( + k_tile_iter, + load_state_k, + smem_pipe_write, + lane_predicate, + k_tile_count_tma, + mcast_mask_b, + iter_to_tile_map); + } + + // Wait for the pipeline MMAs to drain + warpgroup_wait<0>(); + warpgroup_fence_operand(acc_qk); + warpgroup_fence_operand(acc_pv); + + softmax.template step_interleave_begin( + acc_qk, + tiled_mma_qk, + tPcP, + softmax_state, + acc_pv, + tiled_mma_pv, + problem_size, + is_fully_block_sparse, + has_kv_padding, + params.q_shape, + qkv_shape, + params.na_params, + kv_start, + num_tiles, + MultiDimTileShape{}); + + pipeline.consumer_release(smem_pipe_release); + + ++smem_pipe_release; + + pipeline.consumer_wait(smem_pipe_read); + + // MMA PV + auto layout_qk_input = convert_c_layout_to_a_layout( + acc_qk.layout(), + shape<1>(typename decltype(tiled_mma_pv)::LayoutA_TV{})); + + Tensor acc_qk_input = make_tensor(acc_qk.data(), layout_qk_input); + + static_assert(decltype(size<1>(layout_qk_input) == _1{})::value); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<2>(tOrV); i++) { + Tensor acc_qk_element = + make_fragment_like(layout_qk_input(_, _0{}, _0{})); + Tensor acc_qk_element_mk = tensor_op_mk_v(tiled_mma_pv, acc_qk_element); + Tensor acc_qk_input_mk = + tensor_op_mk_v(tiled_mma_pv, acc_qk_input(_, _0{}, i)); + softmax.step_interleave_step(acc_qk_input_mk, softmax_state); + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size(acc_qk_element_mk); j++) { + acc_qk_element_mk(j) = static_cast(acc_qk_input_mk(j)); + } + warpgroup_arrive(); + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<1>(tOrV); j++) { + cute::gemm( + tiled_mma_pv, + acc_qk_element, + tOrV(_, j, i, smem_pipe_read.index()), + acc_pv(_, _0{}, j)); + } + } + warpgroup_commit_batch(); + + // Wait for the pipeline MMAs to drain + pipeline.consumer_release(smem_pipe_release); + ++smem_pipe_release; + + ++smem_pipe_read; + + if (warp_idx == 0) { + load_v.template step_with_iter_tile_map( + k_tile_iter, + load_state_v, + smem_pipe_write, + lane_predicate, + k_tile_count_tma, + mcast_mask_b, + iter_to_tile_map); + } + + tPcP.data() = tPcP.data() + E<1>{} * get<1>(TileShapeQK{}); + } + + k_tile_count += fusion_tile_count; + + CUTLASS_PRAGMA_NO_UNROLL + for (; k_tile_count > 0; --k_tile_count) { + // Allocate QK acc + Tensor acc_qk = + partition_fragment_C(tiled_mma_qk, take<0, 2>(TileShapeQK{})); + + pipeline.consumer_wait(smem_pipe_read); + + // MMA QK + warpgroup_fence_operand(acc_qk); + warpgroup_arrive(); + + gemm_zero_acc( + tiled_mma_qk, + tSrQ(_, _, _, _0{}), + tSrK(_, _, _, smem_pipe_read.index()), + acc_qk); + warpgroup_commit_batch(); + + ++smem_pipe_read; + + if (warp_idx == 0) { + load_k.template step_with_iter_tile_map( + k_tile_iter, + load_state_k, + smem_pipe_write, + lane_predicate, + k_tile_count_tma, + mcast_mask_b, + iter_to_tile_map); + } + + // Wait for the pipeline MMAs to drain + warpgroup_wait<0>(); + warpgroup_fence_operand(acc_qk); + warpgroup_fence_operand(acc_pv); + + softmax.step_interleave_begin( + acc_qk, + tiled_mma_qk, + tPcP, + softmax_state, + acc_pv, + tiled_mma_pv, + problem_size, + is_fully_block_sparse, + has_kv_padding, + params.q_shape, + qkv_shape, + params.na_params, + kv_start, + num_tiles, + MultiDimTileShape{}); + + pipeline.consumer_release(smem_pipe_release); + + ++smem_pipe_release; + + pipeline.consumer_wait(smem_pipe_read); + + // MMA PV + auto layout_qk_input = convert_c_layout_to_a_layout( + acc_qk.layout(), + shape<1>(typename decltype(tiled_mma_pv)::LayoutA_TV{})); + + Tensor acc_qk_input = make_tensor(acc_qk.data(), layout_qk_input); + + static_assert(decltype(size<1>(layout_qk_input) == _1{})::value); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<2>(tOrV); i++) { + Tensor acc_qk_element = + make_fragment_like(layout_qk_input(_, _0{}, _0{})); + Tensor acc_qk_element_mk = tensor_op_mk_v(tiled_mma_pv, acc_qk_element); + Tensor acc_qk_input_mk = + tensor_op_mk_v(tiled_mma_pv, acc_qk_input(_, _0{}, i)); + softmax.step_interleave_step(acc_qk_input_mk, softmax_state); + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size(acc_qk_element_mk); j++) { + acc_qk_element_mk(j) = static_cast(acc_qk_input_mk(j)); + } + warpgroup_arrive(); + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<1>(tOrV); j++) { + cute::gemm( + tiled_mma_pv, + acc_qk_element, + tOrV(_, j, i, smem_pipe_read.index()), + acc_pv(_, _0{}, j)); + } + } + warpgroup_commit_batch(); + + // Wait for the pipeline MMAs to drain + pipeline.consumer_release(smem_pipe_release); + ++smem_pipe_release; + + ++smem_pipe_read; + + if (warp_idx == 0) { + load_v.template step_with_iter_tile_map( + k_tile_iter, + load_state_v, + smem_pipe_write, + lane_predicate, + k_tile_count_tma, + mcast_mask_b, + iter_to_tile_map); + } + + tPcP.data() = tPcP.data() + E<1>{} * get<1>(TileShapeQK{}); + } + + // Wait for the pipeline MMAs to drain + warpgroup_wait<0>(); + warpgroup_fence_operand(acc_pv); + + Tensor lse = softmax.tail(softmax_state, acc_pv, tiled_mma_pv); + + return make_tuple(acc_pv, lse); + } +}; + +} // namespace cutlass::fna::collective diff --git a/natten/csrc/include/natten/cuda/fna_hopper/collective/fna_collective_tma_warpspecialized.hpp b/natten/csrc/include/natten/cuda/fna_hopper/collective/fna_collective_tma_warpspecialized.hpp new file mode 100644 index 00000000..a10fe3ab --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna_hopper/collective/fna_collective_tma_warpspecialized.hpp @@ -0,0 +1,906 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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/gemm/collective/collective_builder.hpp" + +#include "natten/cuda/fna_hopper/collective/fna_collective_load.hpp" +#include "natten/cuda/fna_hopper/collective/fna_collective_softmax.hpp" +#include "natten/cuda/fna_hopper/collective/fna_common.hpp" + +// This is identical to FMHA; reuse +#include "natten/cuda/fmha_hopper/kernel/fmha_options.hpp" + +namespace cutlass::fna::collective { + +using namespace cute; +using cutlass::fmha::kernel::find_option_t; +using cutlass::fmha::kernel::Tag; + +template < + class NADim, + class QTileShape, + class KVTileShape, + class Element_, + class ElementAccumulatorQK_, + class ElementAccumulatorPV_, + class TileShape_, // SeqQ, SeqKV, Head + class LayoutQ_, + class LayoutK_, + class LayoutV_, // SeqX, Head, (Batches) + class Fusion, + class... Options> +struct FnaMainloopTmaWarpSpecializedSm90 { + using Element = Element_; + using ElementAccumulatorQK = ElementAccumulatorQK_; + using ElementAccumulatorPV = ElementAccumulatorPV_; + using TileShape = TileShape_; + + using MultiDimTileShape = cute::tuple; + + static_assert( + size(QTileShape{}) == get<0>(TileShape{}), + "QTileShape doesn't match the size of Q tile in the FMHA kernel."); + static_assert( + size(KVTileShape{}) == get<1>(TileShape{}), + "QTileShape doesn't match the size of Q tile in the FMHA kernel."); + + using LayoutQ = LayoutQ_; + using LayoutK = LayoutK_; + using LayoutV = LayoutV_; + + // Options + static constexpr bool kIsPersistent = + find_option_t::value; + static constexpr bool kIsMainloopLocked = + find_option_t::value; + + static constexpr int NumLoadWarpGroups = 1; + static constexpr int NumMmaWarpGroups = + find_option_t, Options...>::value; + static constexpr int StageCount = + find_option_t, Options...>::value; + static constexpr int StageCountQ = + find_option_t, Options...>::value; + + static const int kOuterLoads = 1; + using StagesQ = cutlass::gemm::collective::StageCount; + using Stages = cutlass::gemm::collective::StageCount; + using ClusterShape = Shape<_1, _1, _1>; + static_assert(StagesQ::value >= NumMmaWarpGroups); + static_assert(Stages::value >= 2); + + // 16B alignment lets us use TMA + static constexpr int Alignment = 16 / sizeof(Element); + + using TileShapeQK = Shape< + decltype(tuple_element_t<0, TileShape>{} / Int{}), + tuple_element_t<1, TileShape>, + tuple_element_t<2, TileShape>>; + + using TileShapePV = decltype(select<0, 2, 1>(TileShapeQK{})); + + using CollectiveMmaQK = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm90, + cutlass::arch::OpClassTensorOp, + Element, + LayoutQ, + Alignment, + Element, + LayoutK, + Alignment, + ElementAccumulatorQK, + TileShapeQK, + ClusterShape, + Stages, + cutlass::gemm::KernelTmaWarpSpecialized>::CollectiveOp; + + using CollectiveMmaPV = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm90, + cutlass::arch::OpClassTensorOp, + // the stride for A does not matter since we do not load from smem at all + Element, + LayoutK, + Alignment, + Element, + decltype(select<1, 0, 2>(LayoutV{})), + Alignment, + ElementAccumulatorPV, + TileShapePV, + ClusterShape, + Stages, + cutlass::gemm::KernelTmaWarpSpecialized>::CollectiveOp; + + using TiledMmaQK = typename CollectiveMmaQK::TiledMma; + using TiledMmaPV = + decltype(convert_to_gmma_rs(typename CollectiveMmaPV::TiledMma{})); + + using SmemLayoutQ = decltype(unstageSmemLayout( + typename CollectiveMmaQK::SmemLayoutA{}, + Int{})); + using SmemLayoutK = typename CollectiveMmaQK::SmemLayoutB; + using SmemLayoutV = typename CollectiveMmaPV::SmemLayoutB; + + using MainloopPipeline = cutlass::PipelineTmaAsync; + using MainloopPipelineQ = cutlass::PipelineTmaAsync; + + using PipelineState = + typename cutlass::PipelineState; + using PipelineStateQ = + typename cutlass::PipelineState; + + static constexpr int kInnerLoadBytes = + size(SmemLayoutK{}(_, _, _0{})) * sizeof(Element); + static constexpr int kOuterLoadBytes = + size(SmemLayoutQ{}(_, _, _0{})) * sizeof(Element); + + using TileShapeOut = TileShapePV; + using TiledMmaOut = TiledMmaPV; + using ElementOut = ElementAccumulatorPV; + + struct SharedStorage { + cute::array_aligned> smem_q; + union { + cute::array_aligned> smem_k; + cute::array_aligned> smem_v; + }; + }; + + struct Arguments { + const Element* ptr_Q; + LayoutQ dQ; + const Element* ptr_K; + LayoutK dK; + const Element* ptr_V; + LayoutV dV; + + // if zero, defaults to 1/sqrt(D) + float scale_softmax = 0.0f; + + // FNA params + NADim q_shape; + NADim kv_shape; + NADim qkv_shape; + NADim window_size; + NADim stride; + NADim dilation; + }; + + 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; + TMA_K tma_load_k; + TMA_V tma_load_v; + + float scale_softmax; + float scale_softmax_log2; + float rp_dropout; + + // FNA params + NADim qkv_shape; + NADim q_shape; + NADim kv_shape; + cute::tuple + na_params; // win, win_left, win_right, stride + bool is_fully_block_sparse; + bool has_kv_padding; + NADim dilation; + bool requires_qkv_fixup; + bool is_dilated; + int num_dilation_groups; + }; + + using LoadQ = cutlass::fna::collective::CollectiveLoadTma< + cutlass::fna::collective::LoadKind::kQ, + MainloopPipelineQ, + Element, + SmemLayoutQ, + TMA_Q>; + + using LoadK = cutlass::fna::collective::CollectiveLoadTma< + cutlass::fna::collective::LoadKind::kK, + MainloopPipeline, + Element, + SmemLayoutK, + TMA_K>; + + using LoadV = cutlass::fna::collective::CollectiveLoadTma< + cutlass::fna::collective::LoadKind::kV, + MainloopPipeline, + Element, + SmemLayoutV, + TMA_V>; + + static_assert( + size(typename CollectiveMmaQK::TiledMma{}) == + size(typename CollectiveMmaPV::TiledMma{})); + + template + static bool can_implement( + ProblemShape const& problem_size, + Arguments const& args) { + return (get<4>(problem_size) <= get<2>(TileShape{})) && + ((get<4>(problem_size) % Alignment) == 0) && + ((get<2>(problem_size) % Alignment) == 0) && + // FNA + evenly_divides(args.q_shape, QTileShape{}) && + evenly_divides(args.kv_shape, KVTileShape{}) && + evenly_divides( + size<0>(problem_size), + size(args.dilation)) && // dilation groups are + // folded into batch + tuple_leq(args.window_size, args.qkv_shape) && + // TODO: check window size * dilation <= qkv shape + tuple_leq(args.stride, args.window_size); + } + + template + static Params to_underlying_arguments( + ProblemShape const& problem_size, + Arguments const& args, + void* workspace) { + auto problem_shape_qk = make_shape( + get<2>(problem_size), + get<3>(problem_size), + get<4>(problem_size), + make_shape(get<0>(problem_size), get<1>(problem_size))); + auto params_qk = CollectiveMmaQK::to_underlying_arguments( + problem_shape_qk, + typename CollectiveMmaQK::Arguments{ + args.ptr_Q, + args.dQ, + args.ptr_K, + args.dK, + }, + /*workspace=*/nullptr); + + auto problem_shape_pv = select<0, 2, 1, 3>(problem_shape_qk); + auto params_pv = CollectiveMmaPV::to_underlying_arguments( + problem_shape_pv, + typename CollectiveMmaPV::Arguments{ + args.ptr_K, + args.dK, // never used, dummy + args.ptr_V, + select<1, 0, 2>(args.dV), + }, + /*workspace=*/nullptr); + + float scale_softmax = args.scale_softmax; + if (scale_softmax == 0.0f) { + scale_softmax = 1.0f / (float)std::sqrt(get<4>(problem_size)); + } + float log2_e = static_cast(std::log2(std::exp(1.0))); + + auto window_left = get_window_left(args.window_size); + auto window_right = get_window_right(args.window_size); + + bool requires_qkv_fixup = not evenly_divides(args.qkv_shape, args.dilation); + + return Params{ + params_qk.tma_load_a, + params_qk.tma_load_b, + params_pv.tma_load_b, + scale_softmax, + log2_e * scale_softmax, + 1.0f, + args.qkv_shape, + args.q_shape, + args.kv_shape, + make_tuple(args.window_size, window_left, window_right, args.stride), + fully_block_sparse( + args.qkv_shape, + args.window_size, + args.stride, + QTileShape{}, + KVTileShape{}), + /* has_kv_padding */ not evenly_divides(args.qkv_shape, KVTileShape{}), + args.dilation, + requires_qkv_fixup, + is_dilated(args.dilation), + size(args.dilation)}; + } + + 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 < + bool kLoadQ, + class BlkCoord, + class ProblemShape, + class LoadWarpBarrier> + CUTLASS_DEVICE void load_kv_maybe_q( + int block_rank_in_cluster, + BlkCoord const& blk_coord, + Params const& params, + ProblemShape const& problem_size, + MainloopPipeline& pipeline, + PipelineState& smem_pipe_write, + MainloopPipelineQ& pipeline_q, + PipelineStateQ& smem_pipe_write_q, + SharedStorage& storage, + LoadWarpBarrier& load_warp_barrier, + bool do_barrier) { + // int fusion_tile_count = + // Fusion{}.get_trip_count(blk_coord, TileShape{}, problem_size); + + auto qkv_shape = params.qkv_shape; + if (params.requires_qkv_fixup) { + qkv_shape = Fusion{}.correct_qkv_shape( + params.qkv_shape, + blk_coord, + params.dilation, + params.num_dilation_groups); + } else if (params.is_dilated) { + qkv_shape = ceil_div(params.qkv_shape, params.dilation); + } + + auto [kv_start, num_tiles] = Fusion{}.get_trip_count( + blk_coord, + MultiDimTileShape{}, + params.q_shape, + qkv_shape, + params.na_params); + + int fusion_tile_count = size(num_tiles); + + auto kv_start_tile = ceil_div(kv_start, KVTileShape{}); + + auto kv_tiled = ceil_div(params.kv_shape, KVTileShape{}); + auto ctr = make_identity_tensor(num_tiles); + auto ctr_offset = domain_offset(kv_start_tile, ctr); + + auto kv_tiled_layout = make_layout(kv_tiled); + + auto iter_to_tile_map = [&ctr_offset, &kv_tiled_layout](int iter) { + return crd2idx(ctr_offset(iter), kv_tiled_layout); + }; + + int lane_predicate = cute::elect_one_sync(); + + uint16_t mcast_mask_b = 0; + + if (lane_predicate == 1) { + if constexpr (cute::is_same_v< + typename CollectiveMmaQK::GmemTiledCopyB, + SM90_TMA_LOAD_MULTICAST>) { + auto block_layout = Layout{}; // (m,n) -> block_id + for (int m = 0; m < size<0>(block_layout); ++m) { + mcast_mask_b |= (uint16_t(1) << block_layout(m, _0{}, Int<0>{})); + } + } + } + + auto q_tile_iter = cute::make_coord_iterator(Int{}); + [[maybe_unused]] int q_tile_count = NumMmaWarpGroups; + + auto k_tile_iter = cute::make_coord_iterator(fusion_tile_count); + int k_tile_count = 2 * fusion_tile_count; + + LoadQ load_q{params.tma_load_q, pipeline_q, storage.smem_q}; + auto load_state_q = load_q.init_state( + _0{}, problem_size, TileShapeQK{}, blk_coord, NumMmaWarpGroups); + + LoadK load_k{params.tma_load_k, pipeline, storage.smem_k}; + auto load_state_k = load_k.init_state( + block_rank_in_cluster, + problem_size, + TileShapeQK{}, + blk_coord, + fusion_tile_count); + + LoadV load_v{params.tma_load_v, pipeline, storage.smem_v}; + auto load_state_v = load_v.init_state( + block_rank_in_cluster, + problem_size, + TileShapePV{}, + blk_coord, + fusion_tile_count); + + if constexpr (kLoadQ) { + load_q.step( + q_tile_iter, + load_state_q, + smem_pipe_write_q, + lane_predicate, + q_tile_count); + } + + load_k.template step_with_iter_tile_map( + k_tile_iter, + load_state_k, + smem_pipe_write, + lane_predicate, + k_tile_count, + mcast_mask_b, + iter_to_tile_map); + + if constexpr (kLoadQ) { + load_q.step( + q_tile_iter, + load_state_q, + smem_pipe_write_q, + lane_predicate, + q_tile_count); + } + + if constexpr (!kLoadQ) { + if (do_barrier) { + load_warp_barrier.arrive(); + load_warp_barrier.wait(/*phase=*/0); + do_barrier = false; + } + } + + load_v.template step_with_iter_tile_map( + k_tile_iter, + load_state_v, + smem_pipe_write, + lane_predicate, + k_tile_count, + mcast_mask_b, + iter_to_tile_map); + + if constexpr (kLoadQ) { + while (q_tile_count > 0) { + load_q.step( + q_tile_iter, + load_state_q, + smem_pipe_write_q, + lane_predicate, + q_tile_count); + } + } + + CUTLASS_PRAGMA_NO_UNROLL + while (k_tile_count > 0) { + load_k.template step_with_iter_tile_map( + k_tile_iter, + load_state_k, + smem_pipe_write, + lane_predicate, + k_tile_count, + mcast_mask_b, + iter_to_tile_map); + load_v.template step_with_iter_tile_map( + k_tile_iter, + load_state_v, + smem_pipe_write, + lane_predicate, + k_tile_count, + mcast_mask_b, + iter_to_tile_map); + } + } + + template + CUTLASS_DEVICE void load_maybe_q( + BlkCoord const& blk_coord, + Params const& params, + ProblemShape const& problem_size, + MainloopPipelineQ& pipeline_q, + PipelineStateQ& smem_pipe_write_q, + SharedStorage& storage, + LoadWarpBarrier& load_warp_barrier, + bool do_barrier) { + int lane_predicate = cute::elect_one_sync(); + + LoadQ load_q{params.tma_load_q, pipeline_q, storage.smem_q}; + auto load_state_q = load_q.init_state( + _0{}, problem_size, TileShapeQK{}, blk_coord, NumMmaWarpGroups); + + auto q_tile_iter = cute::make_coord_iterator(Int{}); + + CUTLASS_PRAGMA_UNROLL + for (int q_tile_count = 0; q_tile_count < NumMmaWarpGroups; + q_tile_count++) { + int count = 1; + load_q.step( + q_tile_iter, load_state_q, smem_pipe_write_q, lane_predicate, count); + if (q_tile_count == 0 && do_barrier) { + load_warp_barrier.arrive(); + load_warp_barrier.wait(/*phase=*/0); + do_barrier = false; + } + } + } + + template < + class BlkCoord, + class ProblemShape, + class MainloopPipelineReducer, + class PipelineStateReducer> + CUTLASS_DEVICE void reduce( + BlkCoord const& blk_coord, + Params const& params, + ProblemShape const& problem_size, + MainloopPipelineReducer& pipeline_reducer, + PipelineStateReducer& smem_pipe_write_reducer, + SharedStorage& storage) { /* no-op */ } + + template < + class BlkCoord, + class ProblemShape, + class MainloopPipelineReducer, + class PipelineStateReducer, + class MathWgOrderBarrier> + CUTLASS_DEVICE auto compute( + BlkCoord const& blk_coord, + BlkCoord const& wg_coord, + Params const& params, + ProblemShape const& problem_size, + MainloopPipeline& pipeline, + PipelineState& smem_pipe_read, + MainloopPipelineQ& pipeline_q, + PipelineStateQ& smem_pipe_read_q, + MainloopPipelineReducer&, + PipelineStateReducer&, + SharedStorage& storage, + MathWgOrderBarrier& math_wg_order_barrier) { + int thread_idx = int(threadIdx.x); + + PipelineState smem_pipe_release = smem_pipe_read; + PipelineStateQ smem_pipe_release_q = smem_pipe_read_q; + + TiledMmaQK tiled_mma_qk; + auto thr_mma_qk = tiled_mma_qk.get_thread_slice(thread_idx); + + // Mainloop setup QK + 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 tSsQ = thr_mma_qk.partition_A(sQ); // (MMA,MMA_M,MMA_K,PIPE) + Tensor tSsK = thr_mma_qk.partition_B(sK); // (MMA,MMA_N,MMA_K,PIPE) + Tensor tSrQ = thr_mma_qk.make_fragment_A(tSsQ); // (MMA,MMA_N,MMA_K,PIPE) + Tensor tSrK = thr_mma_qk.make_fragment_B(tSsK); // (MMA,MMA_M,MMA_N,PIPE) + + // Prepare: MMA PV + TiledMmaPV tiled_mma_pv; + auto thr_mma_pv = tiled_mma_pv.get_thread_slice(thread_idx); + + // Mainloop setup PV + Tensor sV = + make_tensor(make_smem_ptr(storage.smem_v.data()), SmemLayoutV{}); + + Tensor tOsV = thr_mma_pv.partition_B(sV); // (MMA,MMA_N,MMA_K,PIPE) + Tensor tOrV = thr_mma_pv.make_fragment_B(tOsV); // (MMA,MMA_M,MMA_N,PIPE) + + int k_tile_count = 0; + + auto qkv_shape = params.qkv_shape; + bool is_fully_block_sparse = params.is_fully_block_sparse; + bool has_kv_padding = params.has_kv_padding; + if (params.requires_qkv_fixup) { + qkv_shape = Fusion{}.correct_qkv_shape( + params.qkv_shape, + blk_coord, + params.dilation, + params.num_dilation_groups); + is_fully_block_sparse = fully_block_sparse( + qkv_shape, + get<0>(params.na_params), + get<3>(params.na_params), + QTileShape{}, + KVTileShape{}); + has_kv_padding = not evenly_divides(qkv_shape, KVTileShape{}); + } else if (params.is_dilated) { + qkv_shape = ceil_div(params.qkv_shape, params.dilation); + is_fully_block_sparse = fully_block_sparse( + qkv_shape, + get<0>(params.na_params), + get<3>(params.na_params), + QTileShape{}, + KVTileShape{}); + has_kv_padding = not evenly_divides(qkv_shape, KVTileShape{}); + } + + auto [kv_start, num_tiles] = Fusion{}.get_trip_count( + blk_coord, + MultiDimTileShape{}, + params.q_shape, + qkv_shape, + params.na_params); + + int fusion_tile_count = size(num_tiles); + + pipeline_q.consumer_wait(smem_pipe_read_q); + + // mapping into QK accumulator + Tensor cP = make_identity_tensor(take<0, 2>(TileShapeQK{})); + Tensor tPcP = thr_mma_qk.partition_C(cP); + int m_block = get<0>(wg_coord); + tPcP.data() = tPcP.data() + E<0>{} * m_block * get<0>(TileShapeQK{}); + + // Allocate PV acc + Tensor acc_pv = + partition_fragment_C(tiled_mma_pv, take<0, 2>(TileShapePV{})); + + cutlass::fna::collective:: + CollectiveSoftmax + softmax{params}; + auto softmax_state = softmax.init(acc_pv, tiled_mma_pv); + + if (true) { + --k_tile_count; + // Allocate QK acc + Tensor acc_qk = + partition_fragment_C(tiled_mma_qk, take<0, 2>(TileShapeQK{})); + + pipeline.consumer_wait(smem_pipe_read); + math_wg_order_barrier.wait(); + + // MMA QK + warpgroup_fence_operand(acc_qk); + warpgroup_arrive(); + + gemm_zero_acc( + tiled_mma_qk, + tSrQ(_, _, _, smem_pipe_read_q.index()), + tSrK(_, _, _, smem_pipe_read.index()), + acc_qk); + warpgroup_commit_batch(); + math_wg_order_barrier.arrive(); + + ++smem_pipe_read; + + // Wait for the pipeline MMAs to drain + warpgroup_wait<0>(); + warpgroup_fence_operand(acc_qk); + + softmax.step( + acc_qk, + tiled_mma_qk, + tPcP, + softmax_state, + problem_size, + is_fully_block_sparse, + has_kv_padding, + params.q_shape, + qkv_shape, + params.na_params, + kv_start, + num_tiles, + MultiDimTileShape{}); + + Tensor acc_qk_fixed = + make_acc_into_op(acc_qk, typename TiledMmaPV::LayoutA_TV{}); + + pipeline.consumer_wait(smem_pipe_read); + + // MMA PV + warpgroup_fence_operand(acc_pv); + warpgroup_fence_operand(acc_qk_fixed); + warpgroup_arrive(); + + gemm_zero_acc( + tiled_mma_pv, + acc_qk_fixed, + tOrV(_, _, _, smem_pipe_read.index()), + acc_pv); + warpgroup_commit_batch(); + + pipeline.consumer_release(smem_pipe_release); + ++smem_pipe_release; + + // Advance consumer pipeline + ++smem_pipe_read; + tPcP.data() = tPcP.data() + E<1>{} * get<1>(TileShapeQK{}); + } + + CUTLASS_PRAGMA_NO_UNROLL + while (k_tile_count > 0) { + --k_tile_count; + + // Allocate QK acc + Tensor acc_qk = + partition_fragment_C(tiled_mma_qk, take<0, 2>(TileShapeQK{})); + + pipeline.consumer_wait(smem_pipe_read); + + // MMA QK + warpgroup_fence_operand(acc_qk); + warpgroup_arrive(); + + gemm_zero_acc( + tiled_mma_qk, + tSrQ(_, _, _, smem_pipe_read_q.index()), + tSrK(_, _, _, smem_pipe_read.index()), + acc_qk); + warpgroup_commit_batch(); + + ++smem_pipe_read; + auto tok = pipeline.consumer_try_wait(smem_pipe_read); + + // Wait for the pipeline MMAs to drain + warpgroup_wait<0>(); + warpgroup_fence_operand(acc_qk); + warpgroup_fence_operand(acc_pv); + + if constexpr (kIsMainloopLocked) + math_wg_order_barrier.wait(); + softmax.template step( + acc_qk, + tiled_mma_qk, + tPcP, + softmax_state, + acc_pv, + tiled_mma_pv, + problem_size, + is_fully_block_sparse, + has_kv_padding, + params.q_shape, + qkv_shape, + params.na_params, + kv_start, + num_tiles, + MultiDimTileShape{}); + if constexpr (kIsMainloopLocked) + math_wg_order_barrier.arrive(); + + Tensor acc_qk_fixed = + make_acc_into_op(acc_qk, typename TiledMmaPV::LayoutA_TV{}); + + pipeline.consumer_wait(smem_pipe_read, tok); + + // MMA PV + warpgroup_fence_operand(acc_pv); + warpgroup_fence_operand(acc_qk_fixed); + warpgroup_arrive(); + + cute::gemm( + tiled_mma_pv, + acc_qk_fixed, + tOrV(_, _, _, smem_pipe_read.index()), + acc_pv); + warpgroup_commit_batch(); + + pipeline.consumer_release(smem_pipe_release); + ++smem_pipe_release; + + pipeline.consumer_release(smem_pipe_release); + ++smem_pipe_release; + + ++smem_pipe_read; + tPcP.data() = tPcP.data() + E<1>{} * get<1>(TileShapeQK{}); + } + + k_tile_count += fusion_tile_count; + + CUTLASS_PRAGMA_NO_UNROLL + while (k_tile_count > 0) { + --k_tile_count; + + // Allocate QK acc + Tensor acc_qk = + partition_fragment_C(tiled_mma_qk, take<0, 2>(TileShapeQK{})); + + pipeline.consumer_wait(smem_pipe_read); + + // MMA QK + warpgroup_fence_operand(acc_qk); + warpgroup_arrive(); + + gemm_zero_acc( + tiled_mma_qk, + tSrQ(_, _, _, smem_pipe_read_q.index()), + tSrK(_, _, _, smem_pipe_read.index()), + acc_qk); + warpgroup_commit_batch(); + + ++smem_pipe_read; + auto tok = pipeline.consumer_try_wait(smem_pipe_read); + + // Wait for the pipeline MMAs to drain + warpgroup_wait<0>(); + warpgroup_fence_operand(acc_qk); + warpgroup_fence_operand(acc_pv); + + // if constexpr (kIsPersistent) + // if (k_tile_count == 0) + // pipeline_q.consumer_release(smem_pipe_release_q); + + if constexpr (kIsMainloopLocked) + math_wg_order_barrier.wait(); + softmax.step( + acc_qk, + tiled_mma_qk, + tPcP, + softmax_state, + acc_pv, + tiled_mma_pv, + problem_size, + is_fully_block_sparse, + has_kv_padding, + params.q_shape, + qkv_shape, + params.na_params, + kv_start, + num_tiles, + MultiDimTileShape{}); + if constexpr (kIsMainloopLocked) + math_wg_order_barrier.arrive(); + + Tensor acc_qk_fixed = + make_acc_into_op(acc_qk, typename TiledMmaPV::LayoutA_TV{}); + + pipeline.consumer_wait(smem_pipe_read, tok); + + // MMA PV + warpgroup_fence_operand(acc_pv); + warpgroup_fence_operand(acc_qk_fixed); + warpgroup_arrive(); + + cute::gemm( + tiled_mma_pv, + acc_qk_fixed, + tOrV(_, _, _, smem_pipe_read.index()), + acc_pv); + warpgroup_commit_batch(); + + pipeline.consumer_release(smem_pipe_release); + ++smem_pipe_release; + + pipeline.consumer_release(smem_pipe_release); + ++smem_pipe_release; + + ++smem_pipe_read; + tPcP.data() = tPcP.data() + E<1>{} * get<1>(TileShapeQK{}); + } + + if (kIsPersistent) + pipeline_q.consumer_release(smem_pipe_release_q); + + // Wait for the pipeline MMAs to drain + warpgroup_wait<0>(); + warpgroup_fence_operand(acc_pv); + + if (kIsPersistent) + pipeline.consumer_release(smem_pipe_release); + ++smem_pipe_release; + + Tensor lse = softmax.tail(softmax_state, acc_pv, tiled_mma_pv); + + return make_tuple(acc_pv, lse); + } +}; + +} // namespace cutlass::fna::collective diff --git a/natten/csrc/include/natten/cuda/fna_hopper/collective/fna_common.hpp b/natten/csrc/include/natten/cuda/fna_hopper/collective/fna_common.hpp new file mode 100644 index 00000000..5a0d69c5 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna_hopper/collective/fna_common.hpp @@ -0,0 +1,360 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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/kernel_hardware_info.h" + +namespace cutlass::fna::collective { + +using namespace cute; + +template +CUTE_HOST_DEVICE bool tuple_leq(T t0, T t1) { + static_assert(rank(T{}) > 0 && rank(T{}) < 4); + if constexpr (rank(T{}) == 1) { + return get<0>(t0) <= get<0>(t1); + } else if constexpr (rank(T{}) == 2) { + return get<0>(t0) <= get<0>(t1) && get<1>(t0) <= get<1>(t1); + } else { + return get<0>(t0) <= get<0>(t1) && get<1>(t0) <= get<1>(t1) && + get<2>(t0) <= get<2>(t1); + } +} + +template +CUTE_HOST_DEVICE constexpr auto floor_tuple(T0 const& t0, T1 const& t1) { + return transform_leaf( + t0, t1, [&](auto const& a, auto const& b) { return (a / b) * b; }); +} + +template +CUTE_HOST_DEVICE constexpr auto ceil_tuple(T0 const& t0, T1 const& t1) { + return transform_leaf( + t0, t1, [&](auto const& a, auto const& b) { return ceil_div(a, b) * b; }); +} + +template +CUTE_HOST_DEVICE constexpr auto tuple_sub(T0 const& t0, T1 const& t1) { + return transform_leaf( + t0, t1, [&](auto const& a, auto const& b) { return a - b; }); +} + +template +CUTE_HOST_DEVICE constexpr auto tuple_add(T0 const& t0, T1 const& t1) { + return transform_leaf( + t0, t1, [&](auto const& a, auto const& b) { return a + b; }); +} + +template +CUTE_HOST_DEVICE constexpr auto tuple_mul(T0 const& t0, T1 const& t1) { + return transform_leaf( + t0, t1, [&](auto const& a, auto const& b) { return a * b; }); +} + +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; + if constexpr (rA == 2 && rB == 2 && rC == 1) { + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<1>(tA); k_block++) { + cute::gemm(atom, tA(_, k_block), tB(_, k_block), tC); + atom.accumulate_ = GMMA::ScaleOut::One; + } + } else { + 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_ = GMMA::ScaleOut::One; + } + } +} + +template +CUTE_DEVICE void gemm_zero_acc( + Atom& atom, + TA const& tA, + TB const& tB, + TC&& tC) { + atom.accumulate_ = GMMA::ScaleOut::Zero; + gemm_reset_zero_acc(atom, tA, tB, tC); +} + +template +CUTE_DEVICE constexpr typename T::value_type reduce(T const& t, Fn fn) { + if constexpr (decltype(size(t) % _2{} == _0{})::value) { + auto partial = make_tensor(size(t) / _2{}); + CUTE_UNROLL + for (int i = 0; i < size(partial); i++) { + partial(i) = fn(t(i), t(i + size(partial))); + } + return reduce(partial, fn); + } else { + auto result = t(_0{}); + CUTE_UNROLL + for (int i = 1; i < size(t); i++) { + result = fn(result, t(i)); + } + return result; + } +} + +struct fna_max { + CUTE_DEVICE float operator()(float a, float b) { + return ::max(a, b); + } +}; + +template +inline auto __device__ constexpr layout_separate( + Threshold const& thr, + Source const& src, + Reference const& ref) { + auto lt = + filter(transform_layout(src, ref, [&](auto const& s, auto const& r) { + if constexpr (decltype(r < thr)::value) { + return s; + } else { + return make_layout(_1{}, _0{}); + } + })); + auto ge = + filter(transform_layout(src, ref, [&](auto const& s, auto const& r) { + if constexpr (decltype(r >= thr)::value) { + return s; + } else { + return make_layout(_1{}, _0{}); + } + })); + return make_layout(lt, ge); +} + +template +inline auto __device__ constexpr layout_acc_mn( + TiledMma const& tiled_mma, + Acc const& acc) { + auto separated = layout_separate( + get<0>(typename TiledMma::Shape_MNK{}), + get<0>(acc), + stride<1>(typename TiledMma::LayoutC_TV{})); + auto V_M = get<0>(separated); + auto V_N = get<1>(separated); + return make_layout( + make_layout(V_M, get<1>(acc)), make_layout(V_N, get<2>(acc))); +} + +template +inline auto __device__ constexpr layout_op_mk_v( + TiledMma const& tiled_mma, + Acc const& acc) { + return layout_separate( + get<0>(typename TiledMma::Shape_MNK{}), + get<0>(acc), + stride<1>(typename TiledMma::LayoutA_TV{})); +} + +template +inline auto __device__ constexpr tensor_op_mk_v( + TiledMma const& tiled_mma, + Acc&& acc) { + return make_tensor(acc.data(), layout_op_mk_v(tiled_mma, acc.layout())); +} + +template +inline auto __device__ constexpr reduction_target_n(TiledMma const& tiled_mma) { + auto separated = layout_separate( + get<0>(typename TiledMma::Shape_MNK{}), + make_layout(shape<0>(typename TiledMma::LayoutC_TV{})), + stride<0>(typename TiledMma::LayoutC_TV{})); + return get<1>(separated); +} + +template < + template < + cute::GMMA::Major, + cute::GMMA::Major, + cute::GMMA::ScaleIn, + cute::GMMA::ScaleIn> + class Primitive, + cute::GMMA::Major tA, + cute::GMMA::Major tB, + cute::GMMA::ScaleIn sA, + cute::GMMA::ScaleIn sB> +inline auto __device__ constexpr convert_to_gmma_rs( + cute::MMA_Atom> const& tiled_mma) { + using Atom = cute::MMA_Atom>; + using ElementA = typename Atom::ValTypeA; + using ElementB = typename Atom::ValTypeB; + using ElementC = typename Atom::ValTypeC; + using Shape_MNK = typename Atom::Shape_MNK; + using RS = decltype(cute::GMMA::rs_op_selector< + ElementA, + ElementB, + ElementC, + Shape_MNK, + tA, + tB, + sA, + sB>()); + return cute::MMA_Atom{}; +} + +template < + template + class Primitive, + cute::GMMA::ScaleIn sA, + cute::GMMA::ScaleIn sB> +inline auto __device__ constexpr convert_to_gmma_rs( + cute::MMA_Atom> const& tiled_mma) { + using Atom = cute::MMA_Atom>; + using ElementA = typename Atom::ValTypeA; + using ElementB = typename Atom::ValTypeB; + using ElementC = typename Atom::ValTypeC; + using Shape_MNK = typename Atom::Shape_MNK; + constexpr auto tA = cute::GMMA::Major::K; + constexpr auto tB = cute::GMMA::Major::K; + using RS = decltype(cute::GMMA::rs_op_selector< + ElementA, + ElementB, + ElementC, + Shape_MNK, + tA, + tB, + sA, + sB>()); + return cute::MMA_Atom{}; +} + +template +CUTE_DEVICE auto constexpr convert_to_gmma_rs( + cute::TiledMMA const& tiled_mma) { + return cute::TiledMMA{}; +} + +template +CUTE_DEVICE auto constexpr convert_c_layout_to_a_layout( + CLayout const& c, + AValueShape const& a) { + return make_layout( + make_shape(a, shape<1>(c), make_shape(shape<2>(c), size<0>(c) / size(a))), + make_stride( + stride<0>(c), + stride<1>(c), + make_stride(stride<2>(c), size<2>(a) * stride<0, 2>(c)))); +} + +template +CUTE_DEVICE constexpr auto unstageSmemLayout( + Layout const& layout, + Stages stages = {}) { + return composition(layout, make_tuple(_, _, make_layout(stages))); +} + +template +CUTE_DEVICE auto make_acc_into_op( + Accumulator const& acc, + OperandLayout_TV const& operand_layout_tv) { + Tensor operand = make_fragment_like( + convert_c_layout_to_a_layout(acc.layout(), shape<1>(operand_layout_tv))); + Tensor operand_as_acc = make_tensor(operand.data(), acc.layout()); + + cute::copy(acc, operand_as_acc); + + if constexpr (sizeof(Element) == 1) { + // 00 11 22 33 00 11 22 33 acc layout + // 00 00 11 11 22 22 33 33 operand layout + // BB AA AA BB AA BB BB AA conflict-free exchange pattern + // 16-bit exchange; so process two at a time + // potentially + int tid = threadIdx.x % 4; + auto values_u32 = recast(operand); + + CUTE_UNROLL + for (int n = 0; n < size<1>(values_u32); n++) { + CUTE_UNROLL + for (int k = 0; k < size<2>(values_u32); k++) { + CUTE_UNROLL + for (int ii = 0; ii < 8; ii += 4) { + uint32_t values_tmp_0 = values_u32(ii / 2 + 0, n, k); + uint32_t values_tmp_1 = values_u32(ii / 2 + 1, n, k); + + // step A: + // t 1 v 0 -> t 0 v 1 + // t 2 v 0 -> t 1 v 0 + // t 0 v 1 -> t 2 v 0 + // t 3 v 1 -> t 3 v 1 + + int v_to_send = tid == 1 || tid == 2 ? 0 : 1; + int v_to_recv = v_to_send; + int t_to_recv_from = (0x3021 >> (tid * 4)) & 0xF; + + uint32_t values_tmp_a = v_to_send == 0 ? values_tmp_0 : values_tmp_1; + + values_tmp_a = + __shfl_sync(0xFFFFFFFF, values_tmp_a, t_to_recv_from, 4); + + // step B: + // t 0 v 0 -> t 0 v 0 + // t 3 v 0 -> t 1 v 1 + // t 1 v 1 -> t 2 v 1 + // t 2 v 1 -> t 3 v 0 + + v_to_send = 1 - v_to_send; + v_to_recv = 1 - v_to_recv; + t_to_recv_from = (0x2130 >> (tid * 4)) & 0xF; + + uint32_t values_tmp_b = v_to_send == 0 ? values_tmp_0 : values_tmp_1; + + values_tmp_b = + __shfl_sync(0xFFFFFFFF, values_tmp_b, t_to_recv_from, 4); + + values_u32(ii / 2 + 0, n, k) = __byte_perm( + values_tmp_a, values_tmp_b, v_to_send == 0 ? 0x1054 : 0x5410); + values_u32(ii / 2 + 1, n, k) = __byte_perm( + values_tmp_a, values_tmp_b, v_to_send == 0 ? 0x3276 : 0x7632); + } + } + } + } + + return operand; +} + +} // namespace cutlass::fna::collective diff --git a/natten/csrc/include/natten/cuda/fna_hopper/collective/fna_fusion.hpp b/natten/csrc/include/natten/cuda/fna_hopper/collective/fna_fusion.hpp new file mode 100644 index 00000000..003dd2b6 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna_hopper/collective/fna_fusion.hpp @@ -0,0 +1,543 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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" + +#include "natten/cuda/fna_hopper/collective/fna_common.hpp" + +namespace cutlass::fna::collective { + +using namespace cute; + +template +CUTE_HOST_DEVICE int get_win_start( + int index, + int window_left, + int window_right, + int stride, + int length) { + if constexpr (IsCausal) { + int32_t stride_group_leader_idx = + cutlass::fast_min((index / stride) * stride + stride - 1, length - 1); + return cutlass::fast_max( + stride_group_leader_idx - window_left - window_right, 0); + + } else { + auto stride_group_leader_idx = cutlass::fast_min( + ((index / stride) * stride) + (stride / 2), length - 1); + + return cutlass::fast_max(stride_group_leader_idx - window_left, 0) + + ((stride_group_leader_idx + window_right >= length) * + (length - window_right - stride_group_leader_idx - 1)); + } +} + +template +CUTE_HOST_DEVICE int get_win_end( + int index, + int start, + int window_size, + int length) { + if constexpr (IsCausal) { + return cutlass::fast_min(index + 1, length); + } else { + return start + window_size; + } +} + +template +CUTE_HOST_DEVICE auto get_window_start( + Coord index, + NADim window_left, + NADim window_right, + NADim stride, + NADim length) { + static_assert(rank(index) > 0 && rank(index) < 4); + static_assert(rank(index) == rank(Causal{})); + static_assert(rank(index) == rank(length)); + if constexpr (rank(index) == 1) { + return make_tuple(get_win_start(Causal{})>( + get<0>(index), + get<0>(window_left), + get<0>(window_right), + get<0>(stride), + get<0>(length))); + } else if constexpr (rank(index) == 2) { + return make_tuple( + get_win_start(Causal{})>( + get<0>(index), + get<0>(window_left), + get<0>(window_right), + get<0>(stride), + get<0>(length)), + get_win_start(Causal{})>( + get<1>(index), + get<1>(window_left), + get<1>(window_right), + get<1>(stride), + get<1>(length))); + } else { + return make_tuple( + get_win_start(Causal{})>( + get<0>(index), + get<0>(window_left), + get<0>(window_right), + get<0>(stride), + get<0>(length)), + get_win_start(Causal{})>( + get<1>(index), + get<1>(window_left), + get<1>(window_right), + get<1>(stride), + get<1>(length)), + get_win_start(Causal{})>( + get<2>(index), + get<2>(window_left), + get<2>(window_right), + get<2>(stride), + get<2>(length))); + } +} + +template +CUTE_HOST_DEVICE auto get_window_end( + Coord index, + NADim start, + NADim window_size, + NADim length) { + static_assert(rank(index) > 0 && rank(index) < 4); + static_assert(rank(index) == rank(Causal{})); + static_assert(rank(index) == rank(length)); + if constexpr (rank(index) == 1) { + return make_tuple(get_win_end(Causal{})>( + get<0>(index), get<0>(start), get<0>(window_size), get<0>(length))); + } else if constexpr (rank(index) == 2) { + return make_tuple( + get_win_end(Causal{})>( + get<0>(index), get<0>(start), get<0>(window_size), get<0>(length)), + get_win_end(Causal{})>( + get<1>(index), get<1>(start), get<1>(window_size), get<1>(length))); + } else { + return make_tuple( + get_win_end(Causal{})>( + get<0>(index), get<0>(start), get<0>(window_size), get<0>(length)), + get_win_end(Causal{})>( + get<1>(index), get<1>(start), get<1>(window_size), get<1>(length)), + get_win_end(Causal{})>( + get<2>(index), get<2>(start), get<2>(window_size), get<2>(length))); + } +} + +template +CUTE_HOST_DEVICE bool is_neighbor( + Coord kv_coord, + NADim kv_start, + NADim kv_end) { + static_assert(rank(kv_coord) > 0 && rank(kv_coord) < 4); + static_assert(rank(kv_coord) == rank(kv_start)); + if constexpr (rank(kv_coord) == 1) { + return get<0>(kv_coord) >= get<0>(kv_start) && + get<0>(kv_coord) < get<0>(kv_end); + } else if constexpr (rank(kv_coord) == 2) { + return get<0>(kv_coord) >= get<0>(kv_start) && + get<0>(kv_coord) < get<0>(kv_end) && + get<1>(kv_coord) >= get<1>(kv_start) && + get<1>(kv_coord) < get<1>(kv_end); + } else { + return get<0>(kv_coord) >= get<0>(kv_start) && + get<0>(kv_coord) < get<0>(kv_end) && + get<1>(kv_coord) >= get<1>(kv_start) && + get<1>(kv_coord) < get<1>(kv_end) && + get<2>(kv_coord) >= get<2>(kv_start) && + get<2>(kv_coord) < get<2>(kv_end); + } +} + +template +CUTE_HOST_DEVICE bool is_within_bounds(Coord kv_coord, NADim kv_shape) { + static_assert(rank(kv_coord) > 0 && rank(kv_coord) < 4); + static_assert(rank(kv_coord) == rank(kv_shape)); + if constexpr (rank(kv_coord) == 1) { + return get<0>(kv_coord) < get<0>(kv_shape); + } else if constexpr (rank(kv_coord) == 2) { + return get<0>(kv_coord) < get<0>(kv_shape) && + get<1>(kv_coord) < get<1>(kv_shape); + } else { + return get<0>(kv_coord) < get<0>(kv_shape) && + get<1>(kv_coord) < get<1>(kv_shape) && + get<2>(kv_coord) < get<2>(kv_shape); + } +} + +CUTE_HOST_DEVICE +int qkv_fix_dilation(int qkv_shape, int dilation, int dilation_group) { + auto padding = + 1 - ((dilation_group + (dilation - (qkv_shape % dilation))) / dilation); + return (qkv_shape / dilation) + padding; +} + +template +CUTE_HOST_DEVICE auto correct_qkv_shape_wrt_dilation( + NADim qkv_shape, + NADim dilation, + Coord dilation_group) { + static_assert(rank(qkv_shape) > 0 && rank(qkv_shape) < 4); + static_assert(rank(qkv_shape) == rank(dilation_group)); + if constexpr (rank(qkv_shape) == 1) { + return make_tuple(qkv_fix_dilation( + get<0>(qkv_shape), get<0>(dilation), get<0>(dilation_group))); + } else if constexpr (rank(qkv_shape) == 2) { + return make_tuple( + qkv_fix_dilation( + get<0>(qkv_shape), get<0>(dilation), get<0>(dilation_group)), + qkv_fix_dilation( + get<1>(qkv_shape), get<1>(dilation), get<1>(dilation_group))); + } else { + return make_tuple( + qkv_fix_dilation( + get<0>(qkv_shape), get<0>(dilation), get<0>(dilation_group)), + qkv_fix_dilation( + get<1>(qkv_shape), get<1>(dilation), get<1>(dilation_group)), + qkv_fix_dilation( + get<2>(qkv_shape), get<2>(dilation), get<2>(dilation_group))); + } +} + +template +struct NeighborhoodAttentionMaskSm90 { + using Causal = Causal_; + static_assert(rank(Causal{}) >= 1 && rank(Causal{}) < 4); + + // QKV shape Correction + // Only if dilated and input size % dilation != 0 + // NOTE: every warp role has to execute this before using the mask!! + template + CUTLASS_DEVICE auto correct_qkv_shape( + QKVShape const& qkv_shape, // this is pre-padding, pre-token permute, just + // the original shape of the sequence mode in + // the self attention + BlkCoord const& blk_coord, + Dilation const& dilation, + int num_dilation_groups) { + auto batch_idx = get<2, 0>(blk_coord); + + auto dilation_group_idx = batch_idx % num_dilation_groups; + auto dilation_group_crd = idx2crd(dilation_group_idx, dilation); + + return correct_qkv_shape_wrt_dilation( + qkv_shape, dilation, dilation_group_crd); + } + + template < + class BlkCoord, + class MultiDimTileShape, + class QKVShape, + class NAParams> + CUTLASS_DEVICE auto get_trip_count( + BlkCoord const& blk_coord, + MultiDimTileShape const& multi_dim_tile_shapes, + QKVShape const& q_shape, + QKVShape const& qkv_shape, + NAParams const& na_params) { + auto [q_tile_shape, kv_tile_shape] = multi_dim_tile_shapes; + + auto [window_size, window_left, window_right, stride] = na_params; + + auto q_tiled = ceil_div(q_shape, q_tile_shape); + + // Map query index back to the "original" coord + auto q_tile_coord = idx2crd(static_cast(get<0>(blk_coord)), q_tiled); + auto q_coord = tuple_mul(q_tile_coord, q_tile_shape); + + auto q_tile_offset_last = idx2crd(size(q_tile_shape) - 1, q_tile_shape); + auto q_coord_last = tuple_add(q_coord, q_tile_offset_last); + + auto kv_start_actual = get_window_start( + q_coord, window_left, window_right, stride, qkv_shape); + + auto last_kv_start_actual = get_window_start( + q_coord_last, window_left, window_right, stride, qkv_shape); + auto kv_end_actual = get_window_end( + q_coord_last, last_kv_start_actual, window_size, qkv_shape); + + auto kv_start = floor_tuple(kv_start_actual, kv_tile_shape); + auto kv_end = ceil_tuple(kv_end_actual, kv_tile_shape); + + auto kv_diff = tuple_sub(kv_end, kv_start); + auto kv_diff_tiles = ceil_div(kv_diff, kv_tile_shape); + + return make_tuple(kv_start, kv_diff_tiles); + } + + template < + class AccQK, + class IndexQK, + class MultiDimTileShape, + class QKVShape, + class NAParams> + CUTLASS_DEVICE void apply_mask( + AccQK& acc_qk_mn, + IndexQK const& index_qk, + MultiDimTileShape const& multi_dim_tile_shapes, + QKVShape const& q_shape, + QKVShape const& qkv_shape, + NAParams const& na_params, + QKVShape const& blk_kv_offset, + QKVShape const& kv_diff_tiles) { + auto [q_tile_shape, kv_tile_shape] = multi_dim_tile_shapes; + auto [window_size, window_left, window_right, stride] = na_params; + + auto q_tiled = ceil_div(q_shape, q_tile_shape); + + auto [q_idx_first, kv_idx_first] = + index_qk(crd2idx(make_coord(0, 0), acc_qk_mn.layout())); + + // Q coord remap + int q_tile_idx = q_idx_first / size(q_tile_shape); + int q_tile_res = q_idx_first % size(q_tile_shape); + + auto q_tile_coord = idx2crd(q_tile_idx, q_tiled); + auto q_tile_offset = idx2crd(q_tile_res, q_tile_shape); + auto q_thread_offset = + tuple_add(q_tile_offset, tuple_mul(q_tile_coord, q_tile_shape)); + + auto q_ctr = make_identity_tensor(q_tile_shape); + auto q_ctr_offset = domain_offset(q_thread_offset, q_ctr); + + // KV coord remap + int kv_tile_idx = kv_idx_first / size(kv_tile_shape); + int kv_tile_res = kv_idx_first % size(kv_tile_shape); + + auto kv_tile_coord = idx2crd(kv_tile_idx, kv_diff_tiles); + auto kv_tile_offset = idx2crd(kv_tile_res, kv_tile_shape); + auto kv_thread_offset = tuple_add( + kv_tile_offset, + tuple_add(blk_kv_offset, tuple_mul(kv_tile_coord, kv_tile_shape))); + + auto kv_ctr = make_identity_tensor(kv_tile_shape); + auto kv_ctr_offset = domain_offset(kv_thread_offset, kv_ctr); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<0>(acc_qk_mn); i++) { + auto [q_idx, _] = index_qk(crd2idx(make_coord(i, 0), acc_qk_mn.layout())); + + auto q_coord = q_ctr_offset(q_idx - q_idx_first); + + auto kv_start = get_window_start( + q_coord, window_left, window_right, stride, qkv_shape); + auto kv_end = + get_window_end(q_coord, kv_start, window_size, qkv_shape); + + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<1>(acc_qk_mn); j++) { + auto [_, kv_idx] = + index_qk(crd2idx(make_coord(0, j), acc_qk_mn.layout())); + auto kv_coord = kv_ctr_offset(kv_idx - kv_idx_first); + + if (not is_neighbor(kv_coord, kv_start, kv_end)) { + acc_qk_mn(i, j) = -INFINITY; + } + } + } + } + + template < + class AccQK, + class IndexQK, + class MultiDimTileShape, + class QKVShape, + class NAParams> + CUTLASS_DEVICE void apply_padded_mask( + AccQK& acc_qk_mn, + IndexQK const& index_qk, + MultiDimTileShape const& multi_dim_tile_shapes, + QKVShape const& qkv_shape, + NAParams const& na_params, + QKVShape const& blk_kv_offset, + QKVShape const& kv_diff_tiles) { + auto [q_tile_shape, kv_tile_shape] = multi_dim_tile_shapes; + + // KV coord remap + auto [_, kv_idx_first] = + index_qk(crd2idx(make_coord(0, 0), acc_qk_mn.layout())); + int kv_tile_idx = kv_idx_first / size(kv_tile_shape); + int kv_tile_res = kv_idx_first % size(kv_tile_shape); + + auto kv_tile_coord = idx2crd(kv_tile_idx, kv_diff_tiles); + auto kv_tile_offset = idx2crd(kv_tile_res, kv_tile_shape); + auto kv_thread_offset = tuple_add( + kv_tile_offset, + tuple_add(blk_kv_offset, tuple_mul(kv_tile_coord, kv_tile_shape))); + + auto kv_ctr = make_identity_tensor(kv_tile_shape); + auto kv_ctr_offset = domain_offset(kv_thread_offset, kv_ctr); + + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<1>(acc_qk_mn); j++) { + auto [_, kv_idx] = + index_qk(crd2idx(make_coord(0, j), acc_qk_mn.layout())); + + int iter_offset = kv_idx - kv_idx_first; // % size(kv_tile_shape); + auto kv_coord = kv_ctr_offset(iter_offset); + + if (not is_within_bounds(kv_coord, qkv_shape)) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<0>(acc_qk_mn); i++) { + acc_qk_mn(i, j) = -INFINITY; + } + } + } + } +}; + +// Misc +template +CUTE_HOST_DEVICE constexpr auto get_window_left(NADim const& window_size) { + return transform_leaf(window_size, [&](auto const& w) { return w / 2; }); +} + +template +CUTE_HOST_DEVICE constexpr auto get_window_right(NADim const& window_size) { + return transform_leaf( + window_size, [&](auto const& w) { return (w / 2) + ((w % 2) - 1); }); +} + +CUTE_HOST_DEVICE +bool is_fully_block_sparse_int( + int input_size, + int window_size, + int stride, + int q_tile_size, + int kv_tile_size) { + int num_stride_groups = ceil_div(input_size, stride); + int window_left = window_size / 2; + int window_right = (window_size / 2) + ((window_size % 2) - 1); + int stride_group_center = stride / 2; + + int last_stride_group_center = + min((input_size / stride) * stride + stride_group_center, input_size - 1); + int last_stride_group_window_start = get_win_start( + last_stride_group_center, window_left, window_right, stride, input_size); + + return input_size == window_size || + (stride % q_tile_size == 0 && window_size % kv_tile_size == 0 && + (input_size % stride == 0 || + last_stride_group_window_start % kv_tile_size == 0) && + (((window_left - stride_group_center) % kv_tile_size == 0) || + ( + // input_size % stride == 0 && + input_size % kv_tile_size == 0 && num_stride_groups == 2 && + stride + stride_group_center >= input_size - window_size))); +} + +template +CUTE_HOST_DEVICE bool fully_block_sparse( + NADim input_size, + NADim window_size, + NADim stride, + QTile q_tile_shape, + KVTile kv_tile_shape) { +#ifdef NATTEN_DISABLE_FULLY_BLOCK_SPARSE_FAST_PATH + return false; +#else + // Causal masking can never be fully block-sparse (unless tile/block size is + // 1) + // TODO? + if (cute::any_of(Causal{}, [&](auto const& a) { return a; })) { + return false; + } + + static_assert(rank(input_size) > 0 && rank(input_size) < 4); + static_assert(rank(q_tile_shape) == rank(input_size)); + static_assert(rank(kv_tile_shape) == rank(input_size)); + + if constexpr (rank(input_size) == 1) { + return is_fully_block_sparse_int( + get<0>(input_size), + get<0>(window_size), + get<0>(stride), + get<0>(q_tile_shape), + get<0>(kv_tile_shape)); + } else if constexpr (rank(input_size) == 2) { + return is_fully_block_sparse_int( + get<0>(input_size), + get<0>(window_size), + get<0>(stride), + get<0>(q_tile_shape), + get<0>(kv_tile_shape)) && + is_fully_block_sparse_int( + get<1>(input_size), + get<1>(window_size), + get<1>(stride), + get<1>(q_tile_shape), + get<1>(kv_tile_shape)); + } else { + return is_fully_block_sparse_int( + get<0>(input_size), + get<0>(window_size), + get<0>(stride), + get<0>(q_tile_shape), + get<0>(kv_tile_shape)) && + is_fully_block_sparse_int( + get<1>(input_size), + get<1>(window_size), + get<1>(stride), + get<1>(q_tile_shape), + get<1>(kv_tile_shape)) && + is_fully_block_sparse_int( + get<2>(input_size), + get<2>(window_size), + get<2>(stride), + get<2>(q_tile_shape), + get<2>(kv_tile_shape)); + } +#endif +} + +template +CUTE_HOST_DEVICE bool is_dilated(NADim dilation) { + static_assert(rank(dilation) > 0 && rank(dilation) < 4); + + if constexpr (rank(dilation) == 1) { + return get<0>(dilation) != 1; + } else if constexpr (rank(dilation) == 2) { + return get<0>(dilation) != 1 || get<1>(dilation) != 1; + } else { + return get<0>(dilation) != 1 || get<1>(dilation) != 1 || + get<2>(dilation) != 1; + } +} + +} // namespace cutlass::fna::collective diff --git a/natten/csrc/include/natten/cuda/fna_hopper/collective/fna_fusion_bwd.hpp b/natten/csrc/include/natten/cuda/fna_hopper/collective/fna_fusion_bwd.hpp new file mode 100644 index 00000000..b874645d --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna_hopper/collective/fna_fusion_bwd.hpp @@ -0,0 +1,450 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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" + +#include "natten/cuda/fna_hopper/collective/fna_fusion.hpp" + +namespace cutlass::fna::collective { + +using namespace cute; + +// BWD + +// BWD mask +template +CUTE_HOST_DEVICE int get_bwd_win_start( + int index, + int stride_group_offset, + int window_left, + int window_right, + int window_size, + int stride, + int length) { + if constexpr (IsCausal) { + return index; + } else { + auto window_start_pre_trunc = + (index >= window_size) * (index - window_right + stride_group_offset); + return (window_start_pre_trunc / stride) * stride; + } +} + +template +CUTE_HOST_DEVICE int get_bwd_win_end( + int index, + int stride_group_offset, + int window_left, + int window_right, + int window_size, + int stride, + int length) { + if constexpr (IsCausal) { + // Window end is always the last query that attends to this key/value + 1 + // In the strided case, it will always be the last query in a stride + // group + 1, which means it will always be a multiple of stride. + // In the case of incomplete stridegroups (input size % stride != 0), + // we must correct backward window end by incrementing first non-attending + // stridegroup by 1. + auto first_non_attending_stride_group_idx = + (index + window_size) / stride + (index + window_size >= length); + return cutlass::fast_min( + first_non_attending_stride_group_idx * stride, length); + + } else { + auto bound_check = (index >= (length - window_size)); + auto window_end = + ((index + window_left + stride_group_offset + 1) / stride) * stride; + return (bound_check * length) + ((1 - bound_check) * window_end); + } +} + +template +CUTE_HOST_DEVICE auto get_bwd_window_start( + Coord index, + NADim stride_group_offset, + NADim window_left, + NADim window_right, + NADim window_size, + NADim stride, + NADim length) { + static_assert(rank(index) > 0 && rank(index) < 4); + static_assert(rank(index) == rank(Causal{})); + static_assert(rank(index) == rank(length)); + if constexpr (rank(index) == 1) { + return make_tuple(get_bwd_win_start(Causal{})>( + get<0>(index), + get<0>(stride_group_offset), + get<0>(window_left), + get<0>(window_right), + get<0>(window_size), + get<0>(stride), + get<0>(length))); + } else if constexpr (rank(index) == 2) { + return make_tuple( + get_bwd_win_start(Causal{})>( + get<0>(index), + get<0>(stride_group_offset), + get<0>(window_left), + get<0>(window_right), + get<0>(window_size), + get<0>(stride), + get<0>(length)), + get_bwd_win_start(Causal{})>( + get<1>(index), + get<1>(stride_group_offset), + get<1>(window_left), + get<1>(window_right), + get<1>(window_size), + get<1>(stride), + get<1>(length))); + } else { + return make_tuple( + get_bwd_win_start(Causal{})>( + get<0>(index), + get<0>(stride_group_offset), + get<0>(window_left), + get<0>(window_right), + get<0>(window_size), + get<0>(stride), + get<0>(length)), + get_bwd_win_start(Causal{})>( + get<1>(index), + get<1>(stride_group_offset), + get<1>(window_left), + get<1>(window_right), + get<1>(window_size), + get<1>(stride), + get<1>(length)), + get_bwd_win_start(Causal{})>( + get<2>(index), + get<2>(stride_group_offset), + get<2>(window_left), + get<2>(window_right), + get<2>(window_size), + get<2>(stride), + get<2>(length))); + } +} + +template +CUTE_HOST_DEVICE auto get_bwd_window_end( + Coord index, + NADim stride_group_offset, + NADim window_left, + NADim window_right, + NADim window_size, + NADim stride, + NADim length) { + static_assert(rank(index) > 0 && rank(index) < 4); + static_assert(rank(index) == rank(Causal{})); + static_assert(rank(index) == rank(length)); + if constexpr (rank(index) == 1) { + return make_tuple(get_bwd_win_end(Causal{})>( + get<0>(index), + get<0>(stride_group_offset), + get<0>(window_left), + get<0>(window_right), + get<0>(window_size), + get<0>(stride), + get<0>(length))); + } else if constexpr (rank(index) == 2) { + return make_tuple( + get_bwd_win_end(Causal{})>( + get<0>(index), + get<0>(stride_group_offset), + get<0>(window_left), + get<0>(window_right), + get<0>(window_size), + get<0>(stride), + get<0>(length)), + get_bwd_win_end(Causal{})>( + get<1>(index), + get<1>(stride_group_offset), + get<1>(window_left), + get<1>(window_right), + get<1>(window_size), + get<1>(stride), + get<1>(length))); + } else { + return make_tuple( + get_bwd_win_end(Causal{})>( + get<0>(index), + get<0>(stride_group_offset), + get<0>(window_left), + get<0>(window_right), + get<0>(window_size), + get<0>(stride), + get<0>(length)), + get_bwd_win_end(Causal{})>( + get<1>(index), + get<1>(stride_group_offset), + get<1>(window_left), + get<1>(window_right), + get<1>(window_size), + get<1>(stride), + get<1>(length)), + get_bwd_win_end(Causal{})>( + get<2>(index), + get<2>(stride_group_offset), + get<2>(window_left), + get<2>(window_right), + get<2>(window_size), + get<2>(stride), + get<2>(length))); + } +} + +// Backward pass mask +// Does not allow extra KV fusion +template +struct NeighborhoodAttentionBackwardMaskSm90 { + using Causal = Causal_; + static_assert(rank(Causal{}) >= 1 && rank(Causal{}) < 4); + + // Identical to NeighborhoodAttentionMask::correct_qkv_shape + // QKV shape Correction + // Only if dilated and input size % dilation != 0 + // NOTE: every warp role has to execute this before using the mask!! + template + CUTLASS_DEVICE auto correct_qkv_shape( + ProblemShape const& problem_shape, + QKVShape const& qkv_shape, // this is pre-padding, pre-token permute, just + // the original shape of the sequence mode in + // the self attention + BlkCoord const& blk_coord, + Dilation const& dilation, + int num_dilation_groups) { + auto batch_idx = get<2, 0>(blk_coord); + + auto dilation_group_idx = batch_idx % num_dilation_groups; + auto dilation_group_crd = idx2crd(dilation_group_idx, dilation); + + return correct_qkv_shape_wrt_dilation( + qkv_shape, dilation, dilation_group_crd); + } + + // Unlike in forward pass, trip counts are over Q tiles and not KV tiles. + template < + class BlkCoord, + class MultiDimTileShape, + class QKVShape, + class NAParams> + CUTLASS_DEVICE auto get_trip_count( + BlkCoord const& blk_coord, + MultiDimTileShape const& multi_dim_tile_shapes, + QKVShape const& kv_shape, + QKVShape const& qkv_shape, + NAParams const& na_params) { + auto [q_tile_shape, kv_tile_shape] = multi_dim_tile_shapes; + + auto [window_size, window_left, window_right, stride, stride_group_offset] = + na_params; + + auto kv_tiled = ceil_div(kv_shape, kv_tile_shape); + + // Map KV index back to coord + auto kv_tile_coord = idx2crd(static_cast(get<1>(blk_coord)), kv_tiled); + auto kv_coord = tuple_mul(kv_tile_coord, kv_tile_shape); + + auto kv_tile_offset_last = idx2crd(size(kv_tile_shape) - 1, kv_tile_shape); + auto kv_coord_last = tuple_add(kv_coord, kv_tile_offset_last); + + // q start and end instead of kv like in forward pass + auto q_start_actual = get_bwd_window_start( + kv_coord, + stride_group_offset, + window_left, + window_right, + window_size, + stride, + qkv_shape); + + auto last_q_start_actual = get_bwd_window_start( + kv_coord_last, + stride_group_offset, + window_left, + window_right, + window_size, + stride, + qkv_shape); + auto q_end_actual = get_bwd_window_end( + kv_coord_last, + stride_group_offset, + window_left, + window_right, + window_size, + stride, + qkv_shape); + + auto q_start = floor_tuple(q_start_actual, q_tile_shape); + auto q_end = ceil_tuple(q_end_actual, q_tile_shape); + + auto q_diff = tuple_sub(q_end, q_start); + auto q_diff_tiles = ceil_div(q_diff, q_tile_shape); + + return make_tuple(q_start, q_diff_tiles); + } + + template < + class Acc, + class IndexKQ, + class MultiDimTileShape, + class QKVShape, + class NAParams> + CUTLASS_DEVICE void apply_mask( + Acc& acc, + IndexKQ const& index_kq, + MultiDimTileShape const& multi_dim_tile_shapes, + QKVShape const& kv_shape, + QKVShape const& qkv_shape, + NAParams const& na_params, + QKVShape const& blk_q_offset, + QKVShape const& q_diff_tiles) { + auto [q_tile_shape, kv_tile_shape] = multi_dim_tile_shapes; + auto [window_size, window_left, window_right, stride, stride_group_offset] = + na_params; + + auto kv_tiled = ceil_div(kv_shape, kv_tile_shape); + + auto [kv_idx_first, q_idx_first] = index_kq(0); + + // KV coord remap + int kv_tile_idx = kv_idx_first / size(kv_tile_shape); + int kv_tile_res = kv_idx_first % size(kv_tile_shape); + + auto kv_tile_coord = idx2crd(kv_tile_idx, kv_tiled); + auto kv_tile_offset = idx2crd(kv_tile_res, kv_tile_shape); + auto kv_thread_offset = + tuple_add(kv_tile_offset, tuple_mul(kv_tile_coord, kv_tile_shape)); + + auto kv_ctr = make_identity_tensor(kv_tile_shape); + auto kv_ctr_offset = domain_offset(kv_thread_offset, kv_ctr); + + // Q coord remap + int q_tile_idx = q_idx_first / size(q_tile_shape); + int q_tile_res = q_idx_first % size(q_tile_shape); + + auto q_tile_coord = idx2crd(q_tile_idx, q_diff_tiles); + auto q_tile_offset = idx2crd(q_tile_res, q_tile_shape); + auto q_thread_offset = tuple_add( + q_tile_offset, + tuple_add(blk_q_offset, tuple_mul(q_tile_coord, q_tile_shape))); + + auto q_ctr = make_identity_tensor(q_tile_shape); + auto q_ctr_offset = domain_offset(q_thread_offset, q_ctr); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(acc); i++) { + auto [kv_idx, q_idx] = index_kq(i); + + auto q_coord = q_ctr_offset(q_idx - q_idx_first); + auto kv_coord = kv_ctr_offset(kv_idx - kv_idx_first); + + auto q_start = get_bwd_window_start( + kv_coord, + stride_group_offset, + window_left, + window_right, + window_size, + stride, + qkv_shape); + auto q_end = get_bwd_window_end( + kv_coord, + stride_group_offset, + window_left, + window_right, + window_size, + stride, + qkv_shape); + + if (not is_neighbor(q_coord, q_start, q_end)) { + acc(i) = -INFINITY; + } + } + } + + template < + class Acc, + class IndexKQ, + class MultiDimTileShape, + class QKVShape, + class NAParams> + CUTLASS_DEVICE void apply_padded_mask( + Acc& acc, + IndexKQ const& index_kq, + MultiDimTileShape const& multi_dim_tile_shapes, + QKVShape const& qkv_shape, + NAParams const& na_params, + QKVShape const& blk_q_offset, + QKVShape const& q_diff_tiles) { + auto [q_tile_shape, kv_tile_shape] = multi_dim_tile_shapes; + + // Q coord remap + auto [_, q_idx_first] = index_kq(0); + int q_tile_idx = q_idx_first / size(q_tile_shape); + int q_tile_res = q_idx_first % size(q_tile_shape); + + auto q_tile_coord = idx2crd(q_tile_idx, q_diff_tiles); + auto q_tile_offset = idx2crd(q_tile_res, q_tile_shape); + auto q_thread_offset = tuple_add( + q_tile_offset, + tuple_add(blk_q_offset, tuple_mul(q_tile_coord, q_tile_shape))); + + auto q_ctr = make_identity_tensor(q_tile_shape); + auto q_ctr_offset = domain_offset(q_thread_offset, q_ctr); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(acc); i++) { + auto [_, q_idx] = index_kq(i); + + int iter_offset = q_idx - q_idx_first; // % size(kv_tile_shape); + auto q_coord = q_ctr_offset(iter_offset); + + if (not is_within_bounds(q_coord, qkv_shape)) { + acc(i) = -INFINITY; + } + } + } +}; + +// Misc + +template +CUTE_HOST_DEVICE constexpr auto get_bwd_stride_offset(NADim const& stride) { + return transform_leaf( + stride, [&](auto const& s) { return (s - (s / 2) - 1); }); +} + +} // namespace cutlass::fna::collective diff --git a/natten/csrc/include/natten/cuda/fna_hopper/device/fna_bwd_sm90.hpp b/natten/csrc/include/natten/cuda/fna_hopper/device/fna_bwd_sm90.hpp new file mode 100644 index 00000000..f7893e8e --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna_hopper/device/fna_bwd_sm90.hpp @@ -0,0 +1,426 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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 + +/*! + \file + \brief An universal device layer for cutlass 3.x-style kernels. +*/ + +#include +// common +#include "cutlass/cutlass.h" + +#include "natten/cuda/fna_hopper/collective/fna_collective_bwd_tma_warpspecialized.hpp" +#include "natten/cuda/fna_hopper/collective/fna_fusion.hpp" +#include "natten/cuda/fna_hopper/collective/fna_fusion_bwd.hpp" +#include "natten/cuda/fna_hopper/device/fna_sm90.hpp" + +// These are identical to their FMHA counterparts; reuse from FMHA +#include "natten/cuda/fmha_hopper/collective/fmha_epilogue_bwd.hpp" +#include "natten/cuda/fmha_hopper/collective/fmha_varlen.hpp" +#include "natten/cuda/fmha_hopper/kernel/fmha_kernel_bwd_convert.hpp" +#include "natten/cuda/fmha_hopper/kernel/fmha_kernel_bwd_sum_OdO.hpp" +#include "natten/cuda/fmha_hopper/kernel/fmha_kernel_tma_warpspecialized.hpp" +#include "natten/cuda/fmha_hopper/kernel/fmha_tile_scheduler.hpp" + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::fna::device { + +//////////////////////////////////////////////////////////////////////////////// +////////////////////////////// CUTLASS 3.x API ///////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +template < + class ProblemShape, + class Element, + class ElementAccumulator, + class TileShape, + class Fusion, + class QTileShape, + class KVTileShape, + class NADim, + class... Options> +class FnaBwdSm90 { + public: + /// Argument structure: User API + struct Arguments { + ProblemShape problem_size; + + const Element* ptr_Q; + cute::tuple stride_Q; + const Element* ptr_K; + cute::tuple stride_K; + const Element* ptr_V; + cute::tuple stride_V; + + const Element* ptr_O; + cute::tuple stride_O; + const ElementAccumulator* ptr_LSE; + // NATTEN has a different LSE layout + cute::tuple stride_LSE; + + const Element* ptr_dO; + cute::tuple stride_dO; + + Element* ptr_dQ; + cute::tuple stride_dQ; + Element* ptr_dK; + cute::tuple stride_dK; + Element* ptr_dV; + cute::tuple stride_dV; + + // (F)NA parameters + NADim q_shape; + NADim kv_shape; + NADim qkv_shape; + NADim window_size; + NADim stride; + NADim dilation; + + // if zero, defaults to 1/sqrt(D) + float scale_softmax = 0.0f; + + cutlass::KernelHardwareInfo hw_info; + }; + + using OperationSumOdO = cutlass::fna::device::FnaSm90< + cutlass::fmha::kernel:: + FmhaKernelBwdSumOdO>; + using OperationConvert = cutlass::fna::device::FnaSm90< + cutlass::fmha::kernel:: + FmhaKernelBwdConvert>; + + using Mainloop = + cutlass::fna::collective::FnaBwdMainloopTmaWarpSpecializedSm90< + Element, + ElementAccumulator, + TileShape, + Fusion, + QTileShape, + KVTileShape, + NADim, + Options...>; + + static constexpr bool IsVarlen = + cutlass::fmha::collective::is_problem_shape_variable_length( + ProblemShape{}); + using Epilogue = cutlass::fmha::collective::FmhaBwdEpilogueKV< + Element, + ElementAccumulator, + typename Mainloop::TileShapePV, + IsVarlen>; + + using Operation = cutlass::fna::device::FnaSm90< + cutlass::fmha::kernel::FmhaKernelTmaWarpSpecialized< + ProblemShape, + Mainloop, + Epilogue, + cutlass::fmha::kernel::TileSchedulerBwdAdapter< + cutlass::fmha::kernel::IndividualTileScheduler>, + Options...>>; + + struct Params { + OperationSumOdO op_sum_OdO; + Operation op; + OperationConvert op_convert; + ElementAccumulator* dQ_acc; + size_t dQ_acc_size; + }; + + private: + Params params_; + + // Scaled LSE and sum OdO are batch packed, not sequence packed, to avoid + // alignment issues with LSE. + static typename OperationSumOdO::Arguments to_sum_OdO_arguments( + Arguments const& args, + ElementAccumulator* sum_OdO = nullptr, + ElementAccumulator* scaled_lse = nullptr) { + using namespace cute; + auto [B, H, Q_, K, D] = args.problem_size; + D = cutlass::round_up(D, 8); // Alignment + int Q = cutlass::round_up(static_cast(Q_), 8); // Alignment + auto stride_scalar = + make_stride(static_cast(H) * static_cast(Q), Q, _1{}); + // auto log2_e = log2f(expf(1.0f)); + return typename OperationSumOdO::Arguments{ + args.problem_size, + args.ptr_O, + args.stride_O, + args.ptr_dO, + args.stride_dO, + sum_OdO, + stride_scalar, + args.ptr_LSE, + args.stride_LSE, + scaled_lse, + stride_scalar/*, + 1.0f, + log2_e*/}; + } + + // F32 dQ is batch-packed, not sequence-packed + static typename OperationConvert::Arguments to_convert_arguments( + Arguments const& args, + ElementAccumulator* src = nullptr) { + auto [B, H, Q_, K, D] = args.problem_size; + D = cutlass::round_up(D, 8); // Alignment + int Q = cutlass::round_up(static_cast(Q_), 8); // Alignment + auto stride_src_dQ = make_stride( + static_cast(H * D) * static_cast(Q), Q * D, D, _1{}); + return typename OperationConvert::Arguments{ + args.problem_size, + src, + stride_src_dQ, + nullptr, + stride_src_dQ, + nullptr, + stride_src_dQ, + args.ptr_dQ, + args.stride_dQ, + nullptr, + args.stride_dK, + nullptr, + args.stride_dV}; + } + + static typename Operation::Arguments to_bwd_arguments( + Arguments const& args, + ElementAccumulator* sum_OdO = nullptr, + cute::tuple const& stride_sum_OdO = {}, + ElementAccumulator* scaled_lse = nullptr, + cute::tuple const& stride_scaled_lse = {}, + ElementAccumulator* dQ_acc = nullptr, + cute::tuple const& stride_dQ = {}) { + return typename Operation::Arguments{ + args.problem_size, + { + args.ptr_Q, + args.stride_Q, + args.ptr_K, + args.stride_K, + args.ptr_V, + args.stride_V, + args.ptr_dO, + args.stride_dO, + scaled_lse, + stride_scaled_lse, + sum_OdO, + stride_sum_OdO, + dQ_acc, + stride_dQ, + args.scale_softmax, + // FNA args + {args.q_shape, + args.kv_shape, + args.qkv_shape, + args.window_size, + args.stride, + args.dilation}, + }, + {args.ptr_dK, args.stride_dK, args.ptr_dV, args.stride_dV}, + args.hw_info}; + } + + public: + /// Determines whether the GEMM can execute the given problem. + static Status can_implement(Arguments const& args) { + Status status = Status::kSuccess; + + status = OperationSumOdO::can_implement(to_sum_OdO_arguments(args)); + if (status != Status::kSuccess) { + return status; + } + + status = OperationConvert::can_implement(to_convert_arguments(args)); + if (status != Status::kSuccess) { + return status; + } + + status = Operation::can_implement(to_bwd_arguments(args)); + if (status != Status::kSuccess) { + return status; + } + + return status; + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const& args) { + auto [B_, H_, Q_, K, D_] = args.problem_size; + size_t B = static_cast(B_); + size_t H = static_cast(H_); + size_t D = cutlass::round_up(static_cast(D_), 8); // Alignment + size_t Q = cutlass::round_up(static_cast(Q_), 8); // Alignment + size_t workspace_bytes = 0; + // All three intermediary tensors are batch-packed, not sequence packed + // OdO vector + workspace_bytes += B * H * Q * sizeof(ElementAccumulator); + // scaled LSE vector + workspace_bytes += B * H * Q * sizeof(ElementAccumulator); + // FP32 versions of outputs that are churned (start off with Q only) + workspace_bytes += B * H * Q * D * sizeof(ElementAccumulator); + return workspace_bytes; + } + + /// Initializes state from arguments. + Status initialize_split( + Arguments const& args, + void* workspace_dQ, + void* workspace_sum_OdO, + void* workspace_scaled_lse, + cudaStream_t stream = nullptr) { + CUTLASS_TRACE_HOST( + "FnaBwdSm90::initialize_split() - workspace_dQ=" + << workspace_dQ << ", workspace_sum_OdO=" << workspace_sum_OdO + << "stream: " << (stream ? "non-null" : "null")); + + auto [B_, H_, Q_, K, D_] = args.problem_size; + size_t B = static_cast(B_); + size_t H = static_cast(H_); + size_t D = cutlass::round_up(static_cast(D_), 8); // Alignment + size_t Q = cutlass::round_up(static_cast(Q_), 8); // Alignment + ElementAccumulator* sum_OdO = + reinterpret_cast(workspace_sum_OdO); + ElementAccumulator* scaled_lse = + reinterpret_cast(workspace_scaled_lse); + ElementAccumulator* dQ_acc = + reinterpret_cast(workspace_dQ); + params_.dQ_acc = dQ_acc; + params_.dQ_acc_size = B * H * Q * D * sizeof(ElementAccumulator); + auto args_sum_OdO = to_sum_OdO_arguments(args, sum_OdO, scaled_lse); + auto args_convert = to_convert_arguments(args, dQ_acc); + params_.op_sum_OdO.initialize(args_sum_OdO, nullptr, stream); + params_.op_convert.initialize(args_convert, nullptr, stream); + auto args_bwd = to_bwd_arguments( + args, + sum_OdO, + args_sum_OdO.stride_sum_OdO, + scaled_lse, + args_sum_OdO.stride_scaled_lse, + dQ_acc, + args_convert.stride_src_dQ); + params_.op.initialize(args_bwd, nullptr, stream); + + return Status::kSuccess; + } + + /// Initializes state from arguments. + Status initialize( + Arguments const& args, + void* workspace = nullptr, + cudaStream_t stream = nullptr) { + CUTLASS_TRACE_HOST( + "FnaBwdSm90::initialize() - workspace " + << workspace << ", stream: " << (stream ? "non-null" : "null")); + + auto [B_, H_, Q_, K, D_] = args.problem_size; + size_t B = static_cast(B_); + size_t H = static_cast(H_); + size_t D = cutlass::round_up(static_cast(D_), 8); // Alignment + size_t Q = cutlass::round_up(static_cast(Q_), 8); // Alignment + char* workspace_chr = reinterpret_cast(workspace); + ElementAccumulator* sum_OdO = + reinterpret_cast(workspace_chr); + workspace_chr += B * H * Q * sizeof(ElementAccumulator); + ElementAccumulator* scaled_lse = + reinterpret_cast(workspace_chr); + workspace_chr += B * H * Q * sizeof(ElementAccumulator); + ElementAccumulator* dQ_acc = + reinterpret_cast(workspace_chr); + return initialize_split(args, dQ_acc, sum_OdO, scaled_lse, stream); + } + + /// 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("FnaBwdSm90::run()"); + + Status result = Status::kSuccess; + result = params.op_sum_OdO.run(stream); + if (result != Status::kSuccess) { + return result; + } + + auto cuda_result = + cudaMemsetAsync(params.dQ_acc, 0, params.dQ_acc_size, stream); + if (cuda_result != cudaSuccess) { + return Status::kErrorInternal; + } + + result = params.op.run(stream); + if (result != Status::kSuccess) { + return result; + } + + result = params.op_convert.run(stream); + if (result != Status::kSuccess) { + return result; + } + + 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; + } + + /// 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); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::fna::device + +//////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fna_hopper/device/fna_sm90.hpp b/natten/csrc/include/natten/cuda/fna_hopper/device/fna_sm90.hpp new file mode 100644 index 00000000..f12a3971 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna_hopper/device/fna_sm90.hpp @@ -0,0 +1,287 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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 + +#include +// common +#include "cutlass/cutlass.h" + +#if !defined(__CUDACC_RTC__) +#include "cutlass/cluster_launch.hpp" +#include "cutlass/trace.h" +#endif // !defined(__CUDACC_RTC__) + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::fna::device { + +//////////////////////////////////////////////////////////////////////////////// +////////////////////////////// CUTLASS 3.x API ///////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +template +class FnaSm90 { + 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("Universal::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_sm90, + 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_sm90, + 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( + "Universal::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_sm90, + 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("Universal()::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) { + CUTLASS_TRACE_HOST("Universal::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_sm90; + void* kernel_params[] = {¶ms}; + launch_result = ClusterLauncher::launch( + grid, cluster, block, smem_size, stream, kernel, kernel_params); + } else { + launch_result = Status::kSuccess; + cutlass::arch::synclog_setup(); + device_kernel_sm90<<>>(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::fna::device + +//////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/fna_hopper/fna_backward.cuh b/natten/csrc/include/natten/cuda/fna_hopper/fna_backward.cuh new file mode 100644 index 00000000..23c533b2 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna_hopper/fna_backward.cuh @@ -0,0 +1,200 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ + +#pragma once + +#include "cute/tensor.hpp" + +#include "cutlass/cutlass.h" +#include "cutlass/kernel_hardware_info.h" + +#include "natten/cuda/fna_hopper/collective/fna_fusion.hpp" +#include "natten/cuda/fna_hopper/collective/fna_fusion_bwd.hpp" +#include "natten/cuda/fna_hopper/device/fna_bwd_sm90.hpp" + +namespace natten { +namespace cuda { +namespace fna_hopper { + +#if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) + +using namespace cute; +using namespace cutlass::fna; + +template < + typename Element, + class Causal, + class QTileShape, + class KVTileShape, + class TileShape> +struct KernelBackward { + static_assert( + rank(QTileShape{}) == 1 || rank(QTileShape{}) == 2 || + rank(QTileShape{}) == 3); + static_assert(rank(QTileShape{}) == rank(KVTileShape{})); + static_assert(rank(QTileShape{}) == rank(Causal{})); + using NADim = std::conditional_t< + rank(QTileShape{}) == 3, + cute::tuple, + std::conditional_t< + rank(QTileShape{}) == 2, + cute::tuple, + cute::tuple>>; + + using ElementAccumulator = float; + + // B H Q K D + using ProblemShapeType = cute::tuple; + + using Operation = cutlass::fna::device::FnaBwdSm90< + ProblemShapeType, + Element, + ElementAccumulator, + TileShape, + collective::NeighborhoodAttentionBackwardMaskSm90, + QTileShape, + KVTileShape, + NADim>; + + using Arguments = typename Operation::Arguments; + + Operation op; + + Arguments initialize( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch, + int seqlen_Q, + int seqlen_KV, + int heads, + int dim, + NADim q_shape, + NADim kv_shape, + NADim qkv_shape, + NADim window_size, + NADim stride, + NADim dilation, + int device_id, + float attn_scale) { + auto dim_aligned = cutlass::round_up(dim, 8); // alignment + ProblemShapeType problem_shape = ProblemShapeType{ + batch, + heads, + seqlen_Q, + seqlen_KV, + dim_aligned, + }; + + // heads last profile, with torch's "contiguous layout" + // shape: (batch, heads, seqlen, dim) + // stride: (dim*heads*seqlen, dim*heads, dim, 1) + auto stride_Q = make_stride( + static_cast(dim_aligned * heads) * + static_cast(seqlen_Q), + dim_aligned, + dim_aligned * heads, + _1{}); + auto stride_O = stride_Q; + auto stride_K = make_stride( + static_cast(dim_aligned * heads) * + static_cast(seqlen_KV), + dim_aligned, + dim_aligned * heads, + _1{}); + auto stride_V = stride_K; + auto stride_LSE = make_stride( + static_cast(heads) * static_cast(seqlen_Q), + _1{}, + heads); + + cutlass::KernelHardwareInfo hw_info; + hw_info.device_id = device_id; + hw_info.sm_count = + cutlass::KernelHardwareInfo::query_device_multiprocessor_count( + hw_info.device_id); + + Arguments arguments{ + problem_shape, reinterpret_cast(ptr_Q), + stride_Q, reinterpret_cast(ptr_K), + stride_K, reinterpret_cast(ptr_V), + stride_V, reinterpret_cast(ptr_O), + stride_O, reinterpret_cast(ptr_LSE), + stride_LSE, reinterpret_cast(ptr_dO), + stride_O, reinterpret_cast(ptr_dQ), + stride_Q, reinterpret_cast(ptr_dK), + stride_K, reinterpret_cast(ptr_dV), + stride_V, q_shape, + kv_shape, qkv_shape, + window_size, stride, + dilation, attn_scale, + hw_info}; + + return arguments; + } + + size_t get_workspace_size(Arguments const& arguments) { + return Operation::get_workspace_size(arguments); + } + + void run( + Arguments const& arguments, + void* workspace_ptr, + cudaStream_t stream) { + 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; + return; + } + + status = op.initialize(arguments, workspace_ptr, stream); + if (status != cutlass::Status::kSuccess) { + std::cerr + << "Failed to initialize the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(cudaGetLastError()) << std::endl; + return; + } + + // 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; + } + } +}; + +#endif // defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) + +} // namespace fna_hopper +} // namespace cuda +} // namespace natten diff --git a/natten/csrc/include/natten/cuda/fna_hopper/fna_forward.cuh b/natten/csrc/include/natten/cuda/fna_hopper/fna_forward.cuh new file mode 100644 index 00000000..006f8c2d --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna_hopper/fna_forward.cuh @@ -0,0 +1,252 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ + +#pragma once + +#include "cute/tensor.hpp" + +#include "cutlass/cutlass.h" +#include "cutlass/kernel_hardware_info.h" + +#include "natten/cuda/fna_hopper/collective/fna_fusion.hpp" +#include "natten/cuda/fna_hopper/device/fna_sm90.hpp" +#include "natten/cuda/fna_hopper/kernel/fna_kernel_builder.hpp" + +#include "natten/cuda/hopper_fmha_fna.h" + +namespace natten { +namespace cuda { +namespace fna_hopper { + +#if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) + +using namespace cute; +using namespace cutlass::fna::kernel; +using namespace cutlass::fna::collective; +using namespace cutlass::fna; + +template < + typename Element, + class Causal, + class QTileShape, + class KVTileShape, + class TileShape, + natten::cuda::hopper::HopperKernelSchedule KernelSchedule> +struct KernelForward { + static_assert( + rank(QTileShape{}) == 1 || rank(QTileShape{}) == 2 || + rank(QTileShape{}) == 3); + static_assert(rank(QTileShape{}) == rank(KVTileShape{})); + static_assert(rank(QTileShape{}) == rank(Causal{})); + using NADim = std::conditional_t< + rank(QTileShape{}) == 3, + cute::tuple, + std::conditional_t< + rank(QTileShape{}) == 2, + cute::tuple, + cute::tuple>>; + + using ElementAccumulatorQK = float; + using ElementAccumulatorPV = float; + using ElementOut = Element; + + // B H Q K D + using ProblemShapeType = cute::tuple; + + using StrideQ = cute::tuple>; // Q D (B H) + using StrideK = cute::tuple>; // K D (B H) + + using StrideV = + StrideK; // NOTE: StrideV is different for FP8 due to transpose + using StrideO = StrideQ; + + using StrideLSE = cute::tuple>; // Q (B H) + + static_assert( + KernelSchedule == + natten::cuda::hopper::HopperKernelSchedule::NonPersistent || + KernelSchedule == + natten::cuda::hopper::HopperKernelSchedule::WSCooperative || + KernelSchedule == natten::cuda::hopper::HopperKernelSchedule::WSPingpong); + + using DispatchPolicy = std::conditional_t< + KernelSchedule == + natten::cuda::hopper::HopperKernelSchedule::NonPersistent, + cutlass::gemm::KernelTma, + std::conditional_t< + KernelSchedule == + natten::cuda::hopper::HopperKernelSchedule::WSCooperative, + cutlass::gemm::KernelTmaWarpSpecializedCooperative, + cutlass::gemm::KernelTmaWarpSpecializedPingpong>>; + + using Operation = + cutlass::fna::device::FnaSm90, + DispatchPolicy>::Kernel>; + + using Arguments = typename Operation::Arguments; + + Operation op; + + Arguments initialize( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch, + int seqlen_Q, + int seqlen_KV, + int heads, + int dim, + NADim q_shape, + NADim kv_shape, + NADim qkv_shape, + NADim window_size, + NADim stride, + NADim dilation, + int device_id, + float attn_scale) { + auto dim_aligned = cutlass::round_up(dim, 8); // alignment + ProblemShapeType problem_shape = ProblemShapeType{ + batch, + heads, + seqlen_Q, + seqlen_KV, + dim_aligned, + }; + + // heads last profile, with torch's "contiguous layout" + // shape: (batch, heads, seqlen, dim) + // stride: (dim*heads*seqlen, dim*heads, dim, 1) + auto stride_Q = make_stride( + heads * dim_aligned, + _1{}, + make_stride( + static_cast(heads * dim_aligned) * + static_cast(seqlen_Q), + dim_aligned)); + auto stride_O = stride_Q; + auto stride_K = make_stride( + heads * dim_aligned, + _1{}, + make_stride( + static_cast(heads * dim_aligned) * + static_cast(seqlen_KV), + dim_aligned)); + auto stride_V = stride_K; + auto stride_LSE = make_stride( + heads, + make_stride( + static_cast(heads) * static_cast(seqlen_Q), + _1{})); + + cutlass::KernelHardwareInfo hw_info; + hw_info.device_id = device_id; + hw_info.sm_count = + cutlass::KernelHardwareInfo::query_device_multiprocessor_count( + hw_info.device_id); + + Arguments arguments{ + problem_shape, + {reinterpret_cast(ptr_Q), + stride_Q, + reinterpret_cast(ptr_K), + stride_K, + reinterpret_cast(ptr_V), + stride_V, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation}, + {reinterpret_cast(ptr_O), + stride_O, + reinterpret_cast(ptr_LSE), + stride_LSE}, + hw_info}; + + return arguments; + } + + size_t get_workspace_size(Arguments const& arguments) { + return Operation::get_workspace_size(arguments); + } + + void run( + Arguments const& arguments, + void* workspace_ptr, + cudaStream_t stream) { + 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; + return; + } + + status = op.initialize(arguments, workspace_ptr, stream); + if (status != cutlass::Status::kSuccess) { + std::cerr + << "Failed to initialize the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(cudaGetLastError()) << std::endl; + return; + } + + // 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; + } + +#if 0 + cudaError_t result = cudaDeviceSynchronize(); + if (result != cudaSuccess) { + std::cerr << "Error running the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(result) << std::endl; + return; + } +#endif + } +}; + +#endif // defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) + +} // namespace fna_hopper +} // namespace cuda +} // namespace natten diff --git a/natten/csrc/include/natten/cuda/fna_hopper/kernel/fna_kernel_builder.hpp b/natten/csrc/include/natten/cuda/fna_hopper/kernel/fna_kernel_builder.hpp new file mode 100644 index 00000000..b8599a47 --- /dev/null +++ b/natten/csrc/include/natten/cuda/fna_hopper/kernel/fna_kernel_builder.hpp @@ -0,0 +1,232 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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 "natten/cuda/fna_hopper/collective/fna_collective_tma.hpp" +#include "natten/cuda/fna_hopper/collective/fna_collective_tma_warpspecialized.hpp" + +// These are identical to FMHA; reuse +#include "natten/cuda/fmha_hopper/collective/fmha_epilogue.hpp" +#include "natten/cuda/fmha_hopper/collective/fmha_varlen.hpp" +#include "natten/cuda/fmha_hopper/kernel/fmha_kernel_tma.hpp" +#include "natten/cuda/fmha_hopper/kernel/fmha_kernel_tma_warpspecialized.hpp" +#include "natten/cuda/fmha_hopper/kernel/fmha_options.hpp" +#include "natten/cuda/fmha_hopper/kernel/fmha_tile_scheduler.hpp" + +namespace cutlass::fna::kernel { + +using cutlass::fmha::kernel::find_option_t; +using cutlass::fmha::kernel::Option; +using cutlass::fmha::kernel::Tag; + +template < + class NADim, + class QTileShape, + class KVTileShape, + class ProblemShape, + class Element_, + class ElementAccumulatorQK_, + class ElementAccumulatorPV_, + class TileShape_, // BlockQO, BlockKV, BlockHead + class LayoutQ_, + class LayoutK_, + class LayoutV_, + class Fusion, + class DispatchPolicy, + class... Options> +struct FnaBuilder; + +template < + class NADim, + class QTileShape, + class KVTileShape, + class ProblemShape, + class Element, + class ElementAccumulator, + class TileShape, // BlockQO, BlockKV, BlockHead + class Fusion, + class... Options> +struct FnaBuilder< + NADim, + QTileShape, + KVTileShape, + ProblemShape, + Element, + ElementAccumulator, + ElementAccumulator, + TileShape, + cute::tuple>, + cute::tuple>, + cute::tuple>, + Fusion, + cutlass::gemm::KernelTma, + Options...> { + using CollectiveMainloop = cutlass::fna::collective::FnaMainloopTmaSm90< + NADim, + QTileShape, + KVTileShape, + Element, + ElementAccumulator, + TileShape, + Fusion, + Options...>; + + using CollectiveEpilogue = cutlass::fmha::collective::FmhaFwdEpilogueSm90< + Element, + ElementAccumulator, + typename CollectiveMainloop::TileShapePV, + cutlass::fmha::collective::is_problem_shape_variable_length( + ProblemShape{})>; + + using Kernel = cutlass::fmha::kernel::FmhaKernelTma< + ProblemShape, + CollectiveMainloop, + CollectiveEpilogue, + Options...>; +}; + +template < + class NADim, + class QTileShape, + class KVTileShape, + class ProblemShape, + class Element, + class ElementAccumulatorQK, + class ElementAccumulatorPV, + class TileShape, // BlockQO, BlockKV, BlockHead + class LayoutQ, + class LayoutK, + class LayoutV, + class Fusion, + class... Options> +struct FnaBuilder< + NADim, + QTileShape, + KVTileShape, + ProblemShape, + Element, + ElementAccumulatorQK, + ElementAccumulatorPV, + TileShape, + LayoutQ, + LayoutK, + LayoutV, + Fusion, + cutlass::gemm::KernelTmaWarpSpecializedCooperative, + Options...> { + using CollectiveMainloop = + cutlass::fna::collective::FnaMainloopTmaWarpSpecializedSm90< + NADim, + QTileShape, + KVTileShape, + Element, + ElementAccumulatorQK, + ElementAccumulatorPV, + TileShape, + LayoutQ, + LayoutK, + LayoutV, + Fusion, + Options...>; + + using CollectiveEpilogue = cutlass::fmha::collective::FmhaFwdEpilogueSm90< + Element, + ElementAccumulatorPV, + typename CollectiveMainloop::TileShapePV, + cutlass::fmha::collective::is_problem_shape_variable_length( + ProblemShape{})>; + + static constexpr bool kIsPersistent = + find_option_t::value; + using TileScheduler = std::conditional_t< + kIsPersistent, + cutlass::fmha::kernel::PersistentTileScheduler, + cutlass::fmha::kernel::IndividualTileScheduler>; + + using Kernel = cutlass::fmha::kernel::FmhaKernelTmaWarpSpecialized< + ProblemShape, + CollectiveMainloop, + CollectiveEpilogue, + TileScheduler, + Options...>; +}; + +template < + class NADim, + class QTileShape, + class KVTileShape, + class ProblemShape, + class Element, + class ElementAccumulatorQK, + class ElementAccumulatorPV, + class TileShape, // BlockQO, BlockKV, BlockHead + class LayoutQ, + class LayoutK, + class LayoutV, + class Fusion, + class... Options> +struct FnaBuilder< + NADim, + QTileShape, + KVTileShape, + ProblemShape, + Element, + ElementAccumulatorQK, + ElementAccumulatorPV, + TileShape, + LayoutQ, + LayoutK, + LayoutV, + Fusion, + cutlass::gemm::KernelTmaWarpSpecializedPingpong, + Options...> { + using Kernel = typename FnaBuilder< + NADim, + QTileShape, + KVTileShape, + ProblemShape, + Element, + ElementAccumulatorQK, + ElementAccumulatorPV, + TileShape, + LayoutQ, + LayoutK, + LayoutV, + Fusion, + cutlass::gemm::KernelTmaWarpSpecializedCooperative, + Options..., + Option, + Option>::Kernel; +}; + +} // namespace cutlass::fna::kernel diff --git a/natten/csrc/include/natten/cuda/hopper_fmha_fna.h b/natten/csrc/include/natten/cuda/hopper_fmha_fna.h new file mode 100644 index 00000000..d82ab1a5 --- /dev/null +++ b/natten/csrc/include/natten/cuda/hopper_fmha_fna.h @@ -0,0 +1,71 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ + +#pragma once + +#include +#include + +namespace natten { +namespace cuda { +namespace hopper { + +enum class HopperKernelSchedule { + NonPersistent, + WSCooperative, + WSPingpong, + Invalid +}; + +inline const char* to_string(HopperKernelSchedule schedule) { + switch (schedule) { + case HopperKernelSchedule::NonPersistent: + return "NonPersistent"; + case HopperKernelSchedule::WSCooperative: + return "WSCooperative"; + case HopperKernelSchedule::WSPingpong: + return "WSPingpong"; + case HopperKernelSchedule::Invalid: + return "Invalid"; + } + return "Unknown"; +} + +inline auto kernel_type_int_to_enum_type(int kernel_type) { + switch (kernel_type) { + case 0: + return HopperKernelSchedule::NonPersistent; + case 1: + return HopperKernelSchedule::WSCooperative; + case 2: + return HopperKernelSchedule::WSPingpong; + } + throw std::runtime_error( + "Invalid value for argument kernel_type; " + "expected either 0, 1, or 2, got " + + std::to_string(kernel_type)); +} + +} // namespace hopper +} // namespace cuda +} // namespace natten diff --git a/natten/csrc/include/natten/cuda/reduction/compute_delta.cuh b/natten/csrc/include/natten/cuda/reduction/compute_delta.cuh new file mode 100644 index 00000000..20a3ce8a --- /dev/null +++ b/natten/csrc/include/natten/cuda/reduction/compute_delta.cuh @@ -0,0 +1,99 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ + +#pragma once + +#include +#ifdef NATTEN_WITH_CUTLASS +#include +#include +#include +#endif + +namespace natten { +namespace cuda { + +template +void compute_delta( + cudaStream_t stream, + const ElementIn* ptr_O, + const ElementIn* ptr_dO, + ElementOut* ptr_sum_OdO, + int batch, + int heads, + int seqlen_Q, + int dim) { +#ifdef NATTEN_WITH_CUTLASS + using namespace cute; + // B H Q D + using ProblemShape = cute::tuple; + using OperationSumOdO = cutlass::device::DeviceKernel< + cutlass::fmha::kernel:: + FmhaKernelBwdSumOdO>; + + OperationSumOdO op_sum_OdO; + + ProblemShape problem_shape = cute::make_tuple(batch, heads, seqlen_Q, dim); + auto stride_O = make_stride( + static_cast(dim * heads) * static_cast(seqlen_Q), + dim, + dim * heads, + _1{}); + auto stride_dO = stride_O; + auto stride_sum_OdO = make_stride( + static_cast(heads) * static_cast(seqlen_Q), + _1{}, + heads); + + auto args = typename OperationSumOdO::Arguments{ + problem_shape, + ptr_O, + stride_O, + ptr_dO, + stride_dO, + ptr_sum_OdO, + stride_sum_OdO}; + + auto status = OperationSumOdO::can_implement(args); + if (status != cutlass::Status::kSuccess) { + NATTEN_FAILURE( + "`compute_delta` kernel is not supported for this use case."); + } + + status = op_sum_OdO.initialize(args, nullptr, stream); + if (status != cutlass::Status::kSuccess) { + NATTEN_FAILURE("`compute_delta` kernel failed to initialize."); + } + + auto result = op_sum_OdO.run(stream); + if (result != cutlass::Status::kSuccess) { + NATTEN_FAILURE("`compute_delta` kernel launch failed."); + } +#else + NATTEN_FAILURE( + "`compute_delta` is only available when NATTEN is built with CUTLASS."); +#endif +} + +} // namespace cuda +} // namespace natten diff --git a/natten/csrc/include/natten/cuda/reduction/fmha_kernel_bwd_sum_OdO.hpp b/natten/csrc/include/natten/cuda/reduction/fmha_kernel_bwd_sum_OdO.hpp new file mode 100644 index 00000000..3a446c3d --- /dev/null +++ b/natten/csrc/include/natten/cuda/reduction/fmha_kernel_bwd_sum_OdO.hpp @@ -0,0 +1,160 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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/arch/arch.h" +#include "cutlass/cutlass.h" + +namespace cutlass::fmha::kernel { + +using namespace cute; + +template +struct FmhaKernelBwdSumOdO { + static_assert(rank(ProblemShape{}) == 4); + struct Arguments { + ProblemShape problem_shape; + + const Element* ptr_O; + cute::tuple stride_O; + const Element* ptr_dO; + cute::tuple stride_dO; + + ElementAcc* ptr_sum_OdO; + cute::tuple stride_sum_OdO; + }; + + using Params = Arguments; + + using ClusterShape = Shape<_1, _1, _1>; + static constexpr int SharedStorageSize = 0; + + static const int MinBlocksPerMultiprocessor = 1; + static const int MaxThreadsPerBlock = 128; + using ArchTag = cutlass::arch::Sm80; + + 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 const int kBlockQ = 16; + + static const int kNumThreadsD = 8; + static const int kNumThreadsQ = MaxThreadsPerBlock / kNumThreadsD; + static const int kElementsPerLoad = 2; + + static const int kIterationsQ = kBlockQ / kNumThreadsQ; + + static bool can_implement(Arguments const& args) { + return get<3>(args.problem_shape) % kElementsPerLoad == 0; + } + + static dim3 get_grid_shape(Params const& params) { + dim3 grid( + ceil_div(size<2>(params.problem_shape), kBlockQ), + size<1>(params.problem_shape), + size<0>(params.problem_shape)); + return grid; + } + + static dim3 get_block_shape() { + dim3 block(kNumThreadsD, kNumThreadsQ, 1); + return block; + } + + static Params to_underlying_arguments( + Arguments const& args, + void* workspace) { + return args; + } + + CUTLASS_DEVICE void operator()(const Params& params, char* smem) { + auto ptr_O_bh = params.ptr_O + blockIdx.y * get<1>(params.stride_O) + + blockIdx.z * get<0>(params.stride_O); + auto ptr_dO_bh = params.ptr_dO + blockIdx.y * get<1>(params.stride_dO) + + blockIdx.z * get<0>(params.stride_dO); + auto ptr_sum_OdO_bh = params.ptr_sum_OdO + + blockIdx.y * get<1>(params.stride_sum_OdO) + + blockIdx.z * get<0>(params.stride_sum_OdO); + + auto problem_q = get<2>(params.problem_shape); + int seqlen_q = problem_q; + + CUTLASS_PRAGMA_UNROLL + for (int idx_q_t = threadIdx.y; idx_q_t < kBlockQ; + idx_q_t += kNumThreadsQ) { + int idx_q = idx_q_t + kBlockQ * blockIdx.x; + if (idx_q >= seqlen_q) + continue; + ElementAcc acc = 0; + auto ptr_O_bhq = ptr_O_bh + idx_q * get<2>(params.stride_O); + auto ptr_dO_bhq = ptr_dO_bh + idx_q * get<2>(params.stride_dO); + auto ptr_sum_OdO_bhq = + ptr_sum_OdO_bh + idx_q * get<2>(params.stride_sum_OdO); + + for (int idx_d = threadIdx.x * kElementsPerLoad; + idx_d < get<3>(params.problem_shape); + idx_d += kElementsPerLoad * kNumThreadsD) { + Element value_O[kElementsPerLoad]; + Element value_dO[kElementsPerLoad]; + + using Vec = uint_bit_t * kElementsPerLoad>; + *reinterpret_cast(value_O) = + *reinterpret_cast(&ptr_O_bhq[idx_d]); + *reinterpret_cast(value_dO) = + *reinterpret_cast(&ptr_dO_bhq[idx_d]); + + for (int v = 0; v < kElementsPerLoad; v++) { + acc += value_O[v] * value_dO[v]; + } + } + + for (int i = 1; i < kNumThreadsD; i *= 2) { + acc += __shfl_xor_sync((uint32_t)-1, acc, i, kNumThreadsD); + } + + if (threadIdx.x == 0) { + *ptr_sum_OdO_bhq = acc; + } + } + } +}; + +} // namespace cutlass::fmha::kernel diff --git a/natten/csrc/include/natten/cuda/reference/fna_reference_backward.hpp b/natten/csrc/include/natten/cuda/reference/fna_reference_backward.hpp new file mode 100644 index 00000000..91585feb --- /dev/null +++ b/natten/csrc/include/natten/cuda/reference/fna_reference_backward.hpp @@ -0,0 +1,773 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 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 +#include +#include +#include + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace natten { +namespace cuda { +namespace reference { + +template < + int KVTileSize, + int DimPerThread, + class ProblemShape, + class TensorQ, + class TensorK, + class TensorV, + class TensorO, + class TensorLSE, + class TensorDO, + class TensorDQ, /* class TensorDK, class TensorDV, */ + class NADim, + class Causal, + class QKVLayout> +void __global__ fna_bwd_reference_dQ_kernel( + ProblemShape problem_shape, + TensorQ mQ, + TensorK mK, + TensorV mV, + TensorO mO, + TensorLSE mLSE, + TensorDO mDO, + TensorDQ mDQ, /* TensorDK mDK, TensorDV mDV, */ + NADim window_size, + NADim stride, + NADim dilation, + Causal is_causal, + QKVLayout qkv_layout, + float attn_scale, + int num_additional_kv) { + using namespace cute; + + auto attention_mask = + mask::NeighborhoodAttentionReferenceMask( + window_size, stride, dilation, qkv_layout, num_additional_kv); + + using Element = typename TensorO::value_type; + using ElementAccumulator = typename TensorLSE::value_type; + + extern __shared__ char mS_mem[]; + Element* mS = reinterpret_cast(mS_mem); + + cutlass::Array final_acc; + + for (int idx_H = blockIdx.z; idx_H < size<2, 1>(mDQ); idx_H += gridDim.z) { + for (int idx_B = blockIdx.y; idx_B < size<2, 0>(mDQ); idx_B += gridDim.y) { + auto idx_L = make_tuple(idx_B, idx_H); + for (int idx_Q = blockIdx.x; idx_Q < size<0>(mDQ); idx_Q += gridDim.x) { + final_acc.clear(); + auto num_kv_tiles = ceil_div(size<0>(mK), KVTileSize); + for (int kv_tile = 0; kv_tile < num_kv_tiles; ++kv_tile) { + auto kv_tile_offset = kv_tile * KVTileSize; + auto kv_tile_max = kv_tile_offset + KVTileSize; + auto K_max = std::min(kv_tile_max, size<0>(mK)); + auto mS_max = K_max - kv_tile_offset; + + for (int idx_K = kv_tile_offset + threadIdx.x; idx_K < K_max; + idx_K += blockDim.x) { + ElementAccumulator acc_qk = 0; + ElementAccumulator acc_dov = 0; + ElementAccumulator acc_doo = 0; + for (int idx_D0 = 0; idx_D0 < size<1>(mK); idx_D0++) { + acc_qk += mQ(idx_Q, idx_D0, idx_L) * mK(idx_K, idx_D0, idx_L); + } // for idx_D0 + for (int idx_D1 = 0; idx_D1 < size<1>(mV); idx_D1++) { + acc_dov += mDO(idx_Q, idx_D1, idx_L) * mV(idx_K, idx_D1, idx_L); + acc_doo += mDO(idx_Q, idx_D1, idx_L) * mO(idx_Q, idx_D1, idx_L); + } // for idx_D1 + acc_qk *= attn_scale; + acc_dov *= attn_scale; + acc_doo *= attn_scale; + + auto id = make_identity_tensor(make_shape(1, 1)); + auto frag = make_tensor(Shape<_1, _1>{}); + frag(0) = acc_qk; + attention_mask.apply_mask( + frag, + make_tensor( + id.data() + make_arithmetic_tuple(idx_Q, idx_K), + id.layout())); + acc_qk = frag(0); + + mS[idx_K - kv_tile_offset] = static_cast( + exp(acc_qk - mLSE(idx_Q, idx_L)) * (acc_dov - acc_doo)); + } // for idx_K + + __syncthreads(); + + for (int i = 0; i < DimPerThread; ++i) { + int idx_D = threadIdx.x + i * blockDim.x; + if (idx_D < size<1>(mDQ)) { + for (int j = 0; j < mS_max; j++) { + int idx_K = j + kv_tile_offset; + final_acc[i] += mS[j] * mK(idx_K, idx_D, idx_L); + } + } + } // for idx_D + + __syncthreads(); + } + + for (int i = 0; i < DimPerThread; ++i) { + int idx_D = threadIdx.x + i * blockDim.x; + if (idx_D < size<1>(mDQ)) { + mDQ(idx_Q, idx_D, idx_L) = + static_cast(final_acc[i]); + } + } + } + } // idx_B + } // idx_H +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + int QTileSize, + int DimPerThread, + class ProblemShape, + class TensorQ, + class TensorK, + class TensorV, + class TensorO, + class TensorLSE, + class TensorDO, + /* class TensorDQ, */ class TensorDK, /* class TensorDV, */ + class NADim, + class Causal, + class QKVLayout> +void __global__ fna_bwd_reference_dK_kernel( + ProblemShape problem_shape, + TensorQ mQ, + TensorK mK, + TensorV mV, + TensorO mO, + TensorLSE mLSE, + TensorDO mDO, + /* TensorDQ mDQ, */ TensorDK mDK, /* TensorDV mDV, */ + NADim window_size, + NADim stride, + NADim dilation, + Causal is_causal, + QKVLayout qkv_layout, + float attn_scale, + int num_additional_kv) { + using namespace cute; + + auto attention_mask = + mask::NeighborhoodAttentionReferenceMask( + window_size, stride, dilation, qkv_layout, num_additional_kv); + + using Element = typename TensorO::value_type; + using ElementAccumulator = typename TensorLSE::value_type; + + extern __shared__ char mS_mem[]; + Element* mS = reinterpret_cast(mS_mem); + + cutlass::Array final_acc; + + for (int idx_B = blockIdx.y; idx_B < size<2, 0>(mDK); idx_B += gridDim.y) { + for (int idx_K = blockIdx.x; idx_K < size<0>(mDK); idx_K += gridDim.x) { + final_acc.clear(); + for (int idx_H = 0; idx_H < size<2, 1>(mDK); ++idx_H) { + auto idx_L = make_tuple(idx_B, idx_H); + auto num_q_tiles = ceil_div(size<0>(mDO), QTileSize); + for (int q_tile = 0; q_tile < num_q_tiles; ++q_tile) { + auto q_tile_offset = q_tile * QTileSize; + auto q_tile_max = q_tile_offset + QTileSize; + auto Q_max = std::min(q_tile_max, size<0>(mDO)); + auto mS_max = Q_max - q_tile_offset; + + for (int idx_Q = q_tile_offset + threadIdx.x; idx_Q < Q_max; + idx_Q += blockDim.x) { + ElementAccumulator acc_qk = 0; + ElementAccumulator acc_dov = 0; + ElementAccumulator acc_doo = 0; + for (int idx_D0 = 0; idx_D0 < size<1>(mK); idx_D0++) { + acc_qk += mQ(idx_Q, idx_D0, idx_L) * mK(idx_K, idx_D0, idx_L); + } // for idx_D0 + for (int idx_D1 = 0; idx_D1 < size<1>(mV); idx_D1++) { + acc_dov += mDO(idx_Q, idx_D1, idx_L) * mV(idx_K, idx_D1, idx_L); + acc_doo += mDO(idx_Q, idx_D1, idx_L) * mO(idx_Q, idx_D1, idx_L); + } // for idx_D1 + acc_qk *= attn_scale; + acc_dov *= attn_scale; + acc_doo *= attn_scale; + + auto id = make_identity_tensor(make_shape(1, 1)); + auto frag = make_tensor(Shape<_1, _1>{}); + frag(0) = acc_qk; + attention_mask.apply_mask( + frag, + make_tensor( + id.data() + make_arithmetic_tuple(idx_Q, idx_K), + id.layout())); + acc_qk = frag(0); + + mS[idx_Q - q_tile_offset] = static_cast( + exp(acc_qk - mLSE(idx_Q, idx_L)) * (acc_dov - acc_doo)); + } // for idx_Q + + __syncthreads(); + + for (int i = 0; i < DimPerThread; ++i) { + int idx_D = threadIdx.x + i * blockDim.x; + if (idx_D < size<1>(mDK)) { + for (int j = 0; j < mS_max; j++) { + int idx_Q = j + q_tile_offset; + final_acc[i] += mS[j] * mQ(idx_Q, idx_D, idx_L); + } + } + } + + __syncthreads(); + } + + } // for idx_H + + for (int i = 0; i < DimPerThread; ++i) { + int idx_D = threadIdx.x + i * blockDim.x; + if (idx_D < size<1>(mDK)) { + mDK(idx_K, idx_D, make_tuple(idx_B, 0)) = + static_cast(final_acc[i]); + } + } + + } // for idx_K + } // for idx_B +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + int QTileSize, + int DimPerThread, + class ProblemShape, + class TensorQ, + class TensorK, + class TensorV, + class TensorO, + class TensorLSE, + class TensorDO, + /* class TensorDQ, class TensorDK, */ class TensorDV, + class NADim, + class Causal, + class QKVLayout> +void __global__ fna_bwd_reference_dV_kernel( + ProblemShape problem_shape, + TensorQ mQ, + TensorK mK, + TensorV mV, + TensorO mO, + TensorLSE mLSE, + TensorDO mDO, + /* TensorDQ mDQ, TensorDK mDK, */ TensorDV mDV, + NADim window_size, + NADim stride, + NADim dilation, + Causal is_causal, + QKVLayout qkv_layout, + float attn_scale, + int num_additional_kv) { + using namespace cute; + + auto attention_mask = + mask::NeighborhoodAttentionReferenceMask( + window_size, stride, dilation, qkv_layout, num_additional_kv); + + using Element = typename TensorO::value_type; + using ElementAccumulator = typename TensorLSE::value_type; + + extern __shared__ char mS_mem[]; + Element* mS = reinterpret_cast(mS_mem); + + cutlass::Array final_acc; + + for (int idx_B = blockIdx.y; idx_B < size<2, 0>(mDV); idx_B += gridDim.y) { + for (int idx_K = blockIdx.x; idx_K < size<0>(mDV); idx_K += gridDim.x) { + final_acc.clear(); + for (int idx_H = 0; idx_H < size<2, 1>(mDV); ++idx_H) { + auto idx_L = make_tuple(idx_B, idx_H); + auto num_q_tiles = ceil_div(size<0>(mDO), QTileSize); + for (int q_tile = 0; q_tile < num_q_tiles; ++q_tile) { + auto q_tile_offset = q_tile * QTileSize; + auto q_tile_max = q_tile_offset + QTileSize; + auto Q_max = std::min(q_tile_max, size<0>(mDO)); + auto mS_max = Q_max - q_tile_offset; + + for (int idx_Q = q_tile_offset + threadIdx.x; idx_Q < Q_max; + idx_Q += blockDim.x) { + ElementAccumulator acc_qk = 0; + + for (int idx_D0 = 0; idx_D0 < size<1>(mK); idx_D0++) { + ElementAccumulator rQ = mQ(idx_Q, idx_D0, idx_L); + ElementAccumulator rK = mK(idx_K, idx_D0, idx_L); + acc_qk += rQ * rK; + } // for idx_D0 + acc_qk *= attn_scale; + + auto id = make_identity_tensor(make_shape(1, 1)); + auto frag = make_tensor(Shape<_1, _1>{}); + frag(0) = acc_qk; + attention_mask.apply_mask( + frag, + make_tensor( + id.data() + make_arithmetic_tuple(idx_Q, idx_K), + id.layout())); + acc_qk = frag(0); + + mS[idx_Q - q_tile_offset] = + static_cast(exp(acc_qk - mLSE(idx_Q, idx_L))); + } // for idx_Q + + __syncthreads(); + + for (int i = 0; i < DimPerThread; ++i) { + int idx_D = threadIdx.x + i * blockDim.x; + if (idx_D < size<1>(mDV)) { + for (int j = 0; j < mS_max; j++) { + int idx_Q = j + q_tile_offset; + ElementAccumulator rS = mS[/*idx_Q*/ j]; + ElementAccumulator rDO = mDO(idx_Q, idx_D, idx_L); + final_acc[i] += rS * rDO; + } + } + } + + __syncthreads(); + } + } // for idx_H + + for (int i = 0; i < DimPerThread; ++i) { + int idx_D = threadIdx.x + i * blockDim.x; + if (idx_D < size<1>(mDV)) { + mDV(idx_K, idx_D, make_tuple(idx_B, 0)) = + static_cast(final_acc[i]); + } + } + + } // for idx_K + } // for idx_B +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + class ProblemShape, + class TensorQ, + class TensorK, + class TensorV, + class TensorO, + class TensorLSE, + class TensorDO, + /**/ class TensorDQ, /** / class TensorDK, / ** / class TensorDV, / **/ + class NADim, + class Causal, + class QKVLayout> +void fna_bwd_reference_dQ( + ProblemShape problem_shape, + TensorQ mQ, + TensorK mK, + TensorV mV, + TensorO mO, + TensorLSE mLSE, + TensorDO mDO, + /**/ TensorDQ mDQ, /** / TensorDK mDK, / ** / TensorDV mDV, / **/ + NADim window_size, + NADim stride, + NADim dilation, + Causal is_causal, + QKVLayout qkv_layout, + float attn_scale, + int num_additional_kv, + cudaStream_t stream) { + using namespace cute; + + // Only so that we don't oversubscribe shmem when seqlen is large. + static constexpr int KVTileSize = 2048; + static constexpr int MaxDimSupported = 1024; + static constexpr int NumThreads = 256; + static_assert(KVTileSize % NumThreads == 0); + static_assert(MaxDimSupported % NumThreads == 0); + static constexpr int DimPerThread = MaxDimSupported / NumThreads; + + NATTEN_CHECK( + size<1>(mDQ) <= MaxDimSupported, + "Reference kernel only supports up to head dim 1024."); + + dim3 grid(size<0>(mDQ), size<2, 0>(mDQ), size<2, 1>(mDQ)); + dim3 block(NumThreads); + int shared_mem = KVTileSize * sizeof(typename TensorO::value_type); + + NATTEN_CHECK( + shared_mem <= 0xc000, + "Reference kernel oversubscribed shared memory. Please open an issue."); + + fna_bwd_reference_dQ_kernel + <<>>( + problem_shape, + mQ, + mK, + mV, + mO, + mLSE, + mDO, + mDQ, + window_size, + stride, + dilation, + Causal{}, + qkv_layout, + attn_scale, + num_additional_kv); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + class ProblemShape, + class TensorQ, + class TensorK, + class TensorV, + class TensorO, + class TensorLSE, + class TensorDO, + /** / class TensorDQ, / **/ class TensorDK, /** / class TensorDV, / **/ + class NADim, + class Causal, + class QKVLayout> +void fna_bwd_reference_dK( + ProblemShape problem_shape, + TensorQ mQ, + TensorK mK, + TensorV mV, + TensorO mO, + TensorLSE mLSE, + TensorDO mDO, + /** / TensorDQ mDQ, / **/ TensorDK mDK, /** / TensorDV mDV, / **/ + NADim window_size, + NADim stride, + NADim dilation, + Causal is_causal, + QKVLayout qkv_layout, + float attn_scale, + int num_additional_kv, + cudaStream_t stream) { + using namespace cute; + + // Only so that we don't oversubscribe shmem when seqlen is large. + static constexpr int QTileSize = 2048; + static constexpr int MaxDimSupported = 1024; + static constexpr int NumThreads = 256; + static_assert(QTileSize % NumThreads == 0); + static_assert(MaxDimSupported % NumThreads == 0); + static constexpr int DimPerThread = MaxDimSupported / NumThreads; + + NATTEN_CHECK( + size<1>(mDK) <= MaxDimSupported, + "Reference kernel only supports up to head dim 1024."); + + dim3 grid(size<0>(mDK), size<2, 0>(mDK), 1); + dim3 block(NumThreads); + int shared_mem = QTileSize * sizeof(typename TensorO::value_type); + + NATTEN_CHECK( + shared_mem <= 0xc000, + "Reference kernel oversubscribed shared memory. Please open an issue."); + + fna_bwd_reference_dK_kernel + <<>>( + problem_shape, + mQ, + mK, + mV, + mO, + mLSE, + mDO, + mDK, + window_size, + stride, + dilation, + Causal{}, + qkv_layout, + attn_scale, + num_additional_kv); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + class ProblemShape, + class TensorQ, + class TensorK, + class TensorV, + class TensorO, + class TensorLSE, + class TensorDO, + /** / class TensorDQ, / ** / class TensorDK, / **/ class TensorDV, /**/ + class NADim, + class Causal, + class QKVLayout> +void fna_bwd_reference_dV( + ProblemShape problem_shape, + TensorQ mQ, + TensorK mK, + TensorV mV, + TensorO mO, + TensorLSE mLSE, + TensorDO mDO, + /** / TensorDQ mDQ, / ** / TensorDK mDK, / **/ TensorDV mDV, /**/ + NADim window_size, + NADim stride, + NADim dilation, + Causal is_causal, + QKVLayout qkv_layout, + float attn_scale, + int num_additional_kv, + cudaStream_t stream) { + using namespace cute; + + // Only so that we don't oversubscribe shmem when seqlen is large. + static constexpr int QTileSize = 2048; + static constexpr int MaxDimSupported = 1024; + static constexpr int NumThreads = 256; + static_assert(QTileSize % NumThreads == 0); + static_assert(MaxDimSupported % NumThreads == 0); + static constexpr int DimPerThread = MaxDimSupported / NumThreads; + + NATTEN_CHECK( + size<1>(mDV) <= MaxDimSupported, + "Reference kernel only supports up to head dim 1024."); + + dim3 grid(size<0>(mDV), size<2, 0>(mDV), 1); + dim3 block(NumThreads); + int shared_mem = QTileSize * sizeof(typename TensorO::value_type); + + NATTEN_CHECK( + shared_mem <= 0xc000, + "Reference kernel oversubscribed shared memory. Please open an issue."); + + fna_bwd_reference_dV_kernel + <<>>( + problem_shape, + mQ, + mK, + mV, + mO, + mLSE, + mDO, + mDV, + window_size, + stride, + dilation, + Causal{}, + qkv_layout, + attn_scale, + num_additional_kv); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +void fna_reference_backward( + Element* ptr_Q, + Element* ptr_K, + Element* ptr_V, + Element* ptr_O, + Element* ptr_DO, + Element* ptr_DQ, + Element* ptr_DK, + Element* ptr_DV, + ElementLSE* ptr_LSE, + int batch, + int seqlen, + int heads_q, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + NADim qkv_shape, + NADim window_size, + NADim stride, + NADim dilation, + Causal is_causal, + float attn_scale, + cudaStream_t stream) { + using namespace cute; + + auto problem_shape = cute::make_tuple( + seqlen, + seqlen + num_additional_kv, + dim, + cute::make_tuple(cute::make_tuple(heads_kv, batch), heads_q / heads_kv), + dim_value); + + int SQ = size<0>(problem_shape); + int SK = size<1>(problem_shape); + int D = size<2>(problem_shape); + int DV = size<4>(problem_shape); + int H_Q = size<3, 1>(problem_shape); + int H_K = size<3, 0, 0>(problem_shape); + int B = size<3, 0, 1>(problem_shape); + int H = heads_q; + + // heads last profile, with torch's "contiguous layout" + // shape: (batch, seqlen, heads, dim) + // stride: (dim*heads*seqlen, dim*heads, dim, 1) + auto stride_Q = make_stride( + H * D, + _1{}, + make_stride( + make_stride( + H_Q * D, static_cast(H * D) * static_cast(SQ)), + D)); + auto stride_O = make_stride( + H * DV, + _1{}, + make_stride( + make_stride( + H_Q * DV, + static_cast(H * DV) * static_cast(SQ)), + DV)); + auto stride_K = make_stride( + H_K * D, + _1{}, + make_stride( + make_stride( + D, static_cast(H_K * D) * static_cast(SK)), + _0{})); + auto stride_V = make_stride( + H_K * DV, + _1{}, + make_stride( + make_stride( + DV, static_cast(H_K * DV) * static_cast(SK)), + _0{})); + auto stride_LSE = make_stride( + H, + make_stride( + make_stride(H_Q, static_cast(SQ) * static_cast(H)), + _1{})); + + auto mQ = make_tensor( + make_gmem_ptr(ptr_Q), select<0, 2, 3>(problem_shape), stride_Q); + auto mDQ = make_tensor( + make_gmem_ptr(ptr_DQ), select<0, 2, 3>(problem_shape), stride_Q); + + auto mK = make_tensor( + make_gmem_ptr(ptr_K), select<1, 2, 3>(problem_shape), stride_K); + auto mDK = make_tensor( + make_gmem_ptr(ptr_DK), select<1, 2, 3>(problem_shape), stride_K); + + auto mV = make_tensor( + make_gmem_ptr(ptr_V), select<1, 4, 3>(problem_shape), stride_V); + auto mDV = make_tensor( + make_gmem_ptr(ptr_DV), select<1, 4, 3>(problem_shape), stride_V); + + auto mO = make_tensor( + make_gmem_ptr(ptr_O), select<0, 4, 3>(problem_shape), stride_O); + auto mDO = make_tensor( + make_gmem_ptr(ptr_DO), select<0, 4, 3>(problem_shape), stride_O); + + auto mLSE = make_tensor( + make_gmem_ptr(ptr_LSE), select<0, 3>(problem_shape), stride_LSE); + + // For compatibility with pytorch's default contiguous layout + auto qkv_layout = make_layout(qkv_shape, make_qkv_stride(qkv_shape)); + + fna_bwd_reference_dQ( + problem_shape, + mQ, + mK, + mV, + mO, + mLSE, + mDO, + mDQ, + window_size, + stride, + dilation, + Causal{}, + qkv_layout, + attn_scale, + num_additional_kv, + stream); + fna_bwd_reference_dK( + problem_shape, + mQ, + mK, + mV, + mO, + mLSE, + mDO, + mDK, + window_size, + stride, + dilation, + Causal{}, + qkv_layout, + attn_scale, + num_additional_kv, + stream); + fna_bwd_reference_dV( + problem_shape, + mQ, + mK, + mV, + mO, + mLSE, + mDO, + mDV, + window_size, + stride, + dilation, + Causal{}, + qkv_layout, + attn_scale, + num_additional_kv, + stream); +} + +} // namespace reference +} // namespace cuda +} // namespace natten + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/reference/fna_reference_forward.hpp b/natten/csrc/include/natten/cuda/reference/fna_reference_forward.hpp new file mode 100644 index 00000000..3c969f67 --- /dev/null +++ b/natten/csrc/include/natten/cuda/reference/fna_reference_forward.hpp @@ -0,0 +1,329 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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 +#include + +#include +#include +#include +#include + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace natten { +namespace cuda { +namespace reference { + +template < + int KVTileSize, + int DimPerThread, + class ProblemShape, + class TensorQ, + class TensorK, + class TensorV, + class TensorO, + class TensorLSE, + class NADim, + class Causal, + class QKVLayout> +void __global__ fna_reference_kernel( + ProblemShape problem_shape, + TensorQ mQ, + TensorK mK, + TensorV mV, + TensorO mO, + TensorLSE mLSE, + NADim window_size, + NADim stride, + NADim dilation, + Causal is_causal, + QKVLayout qkv_layout, + float attn_scale, + int num_additional_kv) { + using namespace cute; + + auto attention_mask = + mask::NeighborhoodAttentionReferenceMask( + window_size, stride, dilation, qkv_layout, num_additional_kv); + + using Element = typename TensorO::value_type; + using ElementAccumulator = typename TensorLSE::value_type; + + extern __shared__ char mS_mem[]; + ElementAccumulator* mS = reinterpret_cast(mS_mem); + + cutlass::Array final_acc; + + auto id = make_identity_tensor(make_shape(1, 1)); + for (int idx_L = blockIdx.y; idx_L < size<3>(problem_shape); + idx_L += gridDim.y) { + for (int idx_Q = blockIdx.x; idx_Q < size<0>(problem_shape); + idx_Q += gridDim.x) { + final_acc.clear(); + + auto coord_L = idx2crd(idx_L, shape<3>(problem_shape)); + auto coord = cute::make_tuple(idx_Q, _0{}, _0{}, coord_L); + + if (get<0, 0>(coord) >= get<0>(problem_shape)) + continue; + + int offset_Q = 0; + if constexpr (rank<0>(decltype(coord){}) == 2) { + offset_Q = get<0, 1>(coord); + } + + int offset_K = 0; + if constexpr (rank<1>(decltype(coord){}) == 2) { + offset_K = get<1, 1>(coord); + } + + auto num_kv_tiles = ceil_div(size<1>(problem_shape), KVTileSize); + + ElementAccumulator sum = 0; + ElementAccumulator maxS_p = + -std::numeric_limits::infinity(); + ElementAccumulator maxS = + -std::numeric_limits::infinity(); + + for (int kv_tile = 0; kv_tile < num_kv_tiles; ++kv_tile) { + auto kv_tile_offset = kv_tile * KVTileSize; + auto kv_tile_max = kv_tile_offset + KVTileSize; + auto K_max = std::min(kv_tile_max, size<1>(problem_shape)); + auto mS_max = K_max - kv_tile_offset; + + for (int idx_K = kv_tile_offset + threadIdx.x; idx_K < K_max; + idx_K += blockDim.x) { + ElementAccumulator acc = 0; + for (int idx_D = 0; idx_D < size<2>(problem_shape); idx_D++) { + ElementAccumulator eQ = mQ(idx_Q + offset_Q, idx_D, idx_L); + ElementAccumulator eK = mK(idx_K + offset_K, idx_D, idx_L); + acc += eQ * eK; + } + acc = acc * attn_scale; + auto frag = make_tensor(Shape<_1, _1>{}); + frag(0) = acc; + attention_mask.apply_mask( + frag, + make_tensor( + id.data() + make_arithmetic_tuple(idx_Q, idx_K), + id.layout())); + mS[idx_K - kv_tile_offset] = frag(0); + } + + maxS_p = maxS; + __syncthreads(); + + for (int i = 0; i < mS_max; i++) { + maxS = std::max(maxS, mS[i]); + } + + __syncthreads(); + + // If everything so far is masked, just skip to the next tile. + if (maxS == -std::numeric_limits::infinity()) + continue; + + // Scale accumulator so far + if (maxS_p != -std::numeric_limits::infinity()) { + ElementAccumulator correction = expf(maxS_p - maxS); + sum = correction * sum; + for (int i = 0; i < DimPerThread; ++i) { + final_acc[i] = static_cast( + static_cast(final_acc[i]) * correction); + } + } + + for (int idx_K = kv_tile_offset + threadIdx.x; idx_K < K_max; + idx_K += blockDim.x) { + mS[idx_K - kv_tile_offset] = + expf((mS[idx_K - kv_tile_offset] - maxS)); + } + + __syncthreads(); + + for (int i = 0; i < mS_max; i++) { + sum += mS[i]; + } + + for (int i = 0; i < DimPerThread; ++i) { + int idx_D = threadIdx.x + i * blockDim.x; + if (idx_D < size<1>(mO)) { + for (int j = 0; j < mS_max; j++) { + int idx_K = j + kv_tile_offset; + ElementAccumulator eV = mV(idx_K + offset_K, idx_D, idx_L); + ElementAccumulator eK = static_cast(mS[j]); + final_acc[i] += eK * eV; + } + } + } + + __syncthreads(); + } + + for (int i = 0; i < DimPerThread; ++i) { + int idx_D = threadIdx.x + i * blockDim.x; + if (idx_D < size<1>(mO)) { + ElementAccumulator scale = 1.0f / sum; + mO(idx_Q + offset_Q, idx_D, idx_L) = + static_cast(final_acc[i] * scale); + } + } + + if (threadIdx.x == 0) { + mLSE(idx_Q + offset_Q, idx_L) = log(sum) + maxS; + } + } + } +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +void fna_reference_forward( + Element* ptr_Q, + Element* ptr_K, + Element* ptr_V, + Element* ptr_O, + ElementLSE* ptr_LSE, + int batch, + int seqlen, + int heads_q, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + NADim qkv_shape, + NADim window_size, + NADim stride, + NADim dilation, + Causal is_causal, + float attn_scale, + cudaStream_t stream) { + using namespace cute; + + // Only so that we don't oversubscribe shmem when seqlen is large. + static constexpr int KVTileSize = 2048; + static constexpr int MaxDimSupported = 1024; + static constexpr int NumThreads = 256; + static_assert(KVTileSize % NumThreads == 0); + static_assert(MaxDimSupported % NumThreads == 0); + static constexpr int DimPerThread = MaxDimSupported / NumThreads; + + auto problem_shape = cute::make_tuple( + seqlen, + seqlen + num_additional_kv, + dim, + cute::make_tuple(cute::make_tuple(heads_q / heads_kv, heads_kv), batch), + dim_value); + + int SQ = size<0>(problem_shape); + int SK = size<1>(problem_shape); + int D = size<2>(problem_shape); + int DV = size<4>(problem_shape); + int H = size<3, 0>(problem_shape); + int H_K = size<3, 0, 1>(problem_shape); + int H_Q = size<3, 0, 0>(problem_shape); + int B = size<3, 1>(problem_shape); + + // heads last profile, with torch's "contiguous layout" + // shape: (batch, seqlen, heads, dim) + // stride: (dim*heads*seqlen, dim*heads, dim, 1) + auto stride_Q = make_stride( + H * D, _1{}, make_stride(make_stride(D, H_Q * D), H * D * SQ)); + auto stride_O = make_stride( + H * DV, _1{}, make_stride(make_stride(DV, H_Q * DV), H * DV * SQ)); + auto stride_K = make_stride( + H_K * D, _1{}, make_stride(make_stride(_0{}, D), H_K * D * SK)); + auto stride_V = make_stride( + H_K * DV, _1{}, make_stride(make_stride(_0{}, DV), H_K * DV * SK)); + auto stride_LSE = make_stride(H, make_stride(make_stride(_1{}, H_Q), SQ * H)); + + auto mQ = make_tensor( + make_gmem_ptr(ptr_Q), select<0, 2, 3>(problem_shape), stride_Q); + + auto mK = make_tensor( + make_gmem_ptr(ptr_K), select<1, 2, 3>(problem_shape), stride_K); + + auto mV = make_tensor( + make_gmem_ptr(ptr_V), select<1, 4, 3>(problem_shape), stride_V); + + auto mO = make_tensor( + make_gmem_ptr(ptr_O), select<0, 4, 3>(problem_shape), stride_O); + + auto mLSE = make_tensor( + make_gmem_ptr(ptr_LSE), select<0, 3>(problem_shape), stride_LSE); + + // For compatibility with pytorch's default contiguous layout + auto qkv_layout = make_layout(qkv_shape, make_qkv_stride(qkv_shape)); + + NATTEN_CHECK( + size<1>(mO) <= MaxDimSupported, + "Reference kernel only supports up to head dim 1024."); + + dim3 grid(size<0>(mO), size<2>(mO), 1); + dim3 block(NumThreads); + int shared_mem = + KVTileSize * int(sizeof(typename decltype(mLSE)::value_type)); + + NATTEN_CHECK( + shared_mem <= 0xc000, + "Reference kernel oversubscribed shared memory. Please open an issue."); + + fna_reference_kernel + <<>>( + problem_shape, + mQ, + mK, + mV, + mO, + mLSE, + window_size, + stride, + dilation, + Causal{}, + qkv_layout, + attn_scale, + num_additional_kv); + + NATTEN_CUDA_CHECK(cudaGetLastError()); +} + +} // namespace reference +} // namespace cuda +} // namespace natten + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/reference/mask.hpp b/natten/csrc/include/natten/cuda/reference/mask.hpp new file mode 100644 index 00000000..9d38e0cc --- /dev/null +++ b/natten/csrc/include/natten/cuda/reference/mask.hpp @@ -0,0 +1,353 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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 "cutlass/cutlass.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace natten { +namespace cuda { +namespace reference { + +namespace mask { + +using namespace cute; + +template +CUTE_DEVICE constexpr auto get_window_left(NADim const& window_size) { + return transform_leaf(window_size, [&](auto const& w) { return w / 2; }); +} + +template +CUTE_DEVICE constexpr auto get_window_right(NADim const& window_size) { + return transform_leaf( + window_size, [&](auto const& w) { return (w / 2) + ((w % 2) - 1); }); +} + +template +CUTE_DEVICE int get_win_start( + int index, + int window_left, + int window_right, + int stride, + int length) { + if constexpr (IsCausal) { + int32_t stride_group_leader_idx = + cutlass::fast_min((index / stride) * stride + stride - 1, length - 1); + return cutlass::fast_max( + stride_group_leader_idx - window_left - window_right, 0); + + } else { + auto stride_group_leader_idx = cutlass::fast_min( + ((index / stride) * stride) + (stride / 2), length - 1); + + return cutlass::fast_max(stride_group_leader_idx - window_left, 0) + + ((stride_group_leader_idx + window_right >= length) * + (length - window_right - stride_group_leader_idx - 1)); + } +} + +template +CUTE_DEVICE int get_win_end(int index, int start, int window_size, int length) { + if constexpr (IsCausal) { + return cutlass::fast_min(index + 1, length); + } else { + return start + window_size; + } +} + +template +CUTE_DEVICE auto get_window_start( + NADim index, + NADim window_left, + NADim window_right, + NADim stride, + NADim length) { + static_assert(rank(index) > 0 && rank(index) < 4); + static_assert(rank(index) == rank(Causal{})); + if constexpr (rank(index) == 1) { + return make_tuple(get_win_start(Causal{})>( + get<0>(index), + get<0>(window_left), + get<0>(window_right), + get<0>(stride), + get<0>(length))); + } else if constexpr (rank(index) == 2) { + return make_tuple( + get_win_start(Causal{})>( + get<0>(index), + get<0>(window_left), + get<0>(window_right), + get<0>(stride), + get<0>(length)), + get_win_start(Causal{})>( + get<1>(index), + get<1>(window_left), + get<1>(window_right), + get<1>(stride), + get<1>(length))); + } else { + return make_tuple( + get_win_start(Causal{})>( + get<0>(index), + get<0>(window_left), + get<0>(window_right), + get<0>(stride), + get<0>(length)), + get_win_start(Causal{})>( + get<1>(index), + get<1>(window_left), + get<1>(window_right), + get<1>(stride), + get<1>(length)), + get_win_start(Causal{})>( + get<2>(index), + get<2>(window_left), + get<2>(window_right), + get<2>(stride), + get<2>(length))); + } +} + +template +CUTE_DEVICE auto get_window_end( + NADim index, + NADim start, + NADim window_size, + NADim length) { + static_assert(rank(index) > 0 && rank(index) < 4); + static_assert(rank(index) == rank(Causal{})); + if constexpr (rank(index) == 1) { + return make_tuple(get_win_end(Causal{})>( + get<0>(index), get<0>(start), get<0>(window_size), get<0>(length))); + } else if constexpr (rank(index) == 2) { + return make_tuple( + get_win_end(Causal{})>( + get<0>(index), get<0>(start), get<0>(window_size), get<0>(length)), + get_win_end(Causal{})>( + get<1>(index), get<1>(start), get<1>(window_size), get<1>(length))); + } else { + return make_tuple( + get_win_end(Causal{})>( + get<0>(index), get<0>(start), get<0>(window_size), get<0>(length)), + get_win_end(Causal{})>( + get<1>(index), get<1>(start), get<1>(window_size), get<1>(length)), + get_win_end(Causal{})>( + get<2>(index), get<2>(start), get<2>(window_size), get<2>(length))); + } +} + +template +CUTE_DEVICE bool is_neighbor(Coord kv_coord, NADim kv_start, NADim kv_end) { + static_assert(rank(kv_coord) > 0 && rank(kv_coord) < 4); + static_assert(rank(kv_coord) == rank(kv_start)); + if constexpr (rank(kv_coord) == 1) { + return get<0>(kv_coord) >= get<0>(kv_start) && + get<0>(kv_coord) < get<0>(kv_end); + } else if constexpr (rank(kv_coord) == 2) { + return get<0>(kv_coord) >= get<0>(kv_start) && + get<0>(kv_coord) < get<0>(kv_end) && + get<1>(kv_coord) >= get<1>(kv_start) && + get<1>(kv_coord) < get<1>(kv_end); + } else { + return get<0>(kv_coord) >= get<0>(kv_start) && + get<0>(kv_coord) < get<0>(kv_end) && + get<1>(kv_coord) >= get<1>(kv_start) && + get<1>(kv_coord) < get<1>(kv_end) && + get<2>(kv_coord) >= get<2>(kv_start) && + get<2>(kv_coord) < get<2>(kv_end); + } +} + +template +CUTE_DEVICE constexpr auto floor_div_tuple(T0 const& t0, T1 const& t1) { + return transform_leaf( + t0, t1, [&](auto const& a, auto const& b) { return a / b; }); +} + +template +CUTE_DEVICE constexpr auto mod_tuple(T0 const& t0, T1 const& t1) { + return transform_leaf( + t0, t1, [&](auto const& a, auto const& b) { return a % b; }); +} + +template +CUTE_DEVICE bool is_equal(Coord a, Coord b) { + static_assert(rank(a) > 0 && rank(a) < 4); + static_assert(rank(a) == rank(b)); + if constexpr (rank(a) == 1) { + return get<0>(a) == get<0>(b); + } else if constexpr (rank(a) == 2) { + return get<0>(a) == get<0>(b) && get<1>(a) == get<1>(b); + } else { + return get<0>(a) == get<0>(b) && get<1>(a) == get<1>(b) && + get<2>(a) == get<2>(b); + } +} + +CUTE_HOST_DEVICE +int qkv_fix_dilation(int qkv_shape, int dilation, int dilation_group) { + auto padding = + 1 - ((dilation_group + (dilation - (qkv_shape % dilation))) / dilation); + return (qkv_shape / dilation) + padding; +} + +template +CUTE_HOST_DEVICE auto correct_qkv_shape_wrt_dilation( + NADim qkv_shape, + NADim dilation, + Coord dilation_group) { + static_assert(rank(qkv_shape) > 0 && rank(qkv_shape) < 4); + static_assert(rank(qkv_shape) == rank(dilation_group)); + if constexpr (rank(qkv_shape) == 1) { + return make_tuple(qkv_fix_dilation( + get<0>(qkv_shape), get<0>(dilation), get<0>(dilation_group))); + } else if constexpr (rank(qkv_shape) == 2) { + return make_tuple( + qkv_fix_dilation( + get<0>(qkv_shape), get<0>(dilation), get<0>(dilation_group)), + qkv_fix_dilation( + get<1>(qkv_shape), get<1>(dilation), get<1>(dilation_group))); + } else { + return make_tuple( + qkv_fix_dilation( + get<0>(qkv_shape), get<0>(dilation), get<0>(dilation_group)), + qkv_fix_dilation( + get<1>(qkv_shape), get<1>(dilation), get<1>(dilation_group)), + qkv_fix_dilation( + get<2>(qkv_shape), get<2>(dilation), get<2>(dilation_group))); + } +} + +// Not using CuTe tilers here because we'd need to make dilation a layout +// as well since CuTe and torch identity layouts aren't the same, and +// that creeps into weird places that we just don't want for a reference +// kernel. +template +CUTE_DEVICE auto map_idx_to_di_coords( + int idx, + NADim dilation, + QKVLayout qkv_layout) { + auto coord_global = idx2crd(idx, qkv_layout.shape(), qkv_layout.stride()); + + auto coord_within_di_group = floor_div_tuple(coord_global, dilation); + auto dilation_group_crd = mod_tuple(coord_global, dilation); + + return cute::make_tuple(coord_within_di_group, dilation_group_crd); +} + +// Reference Mask +// Handles all parameters (window size, stride, causal, AND dilation) +// in the same place. +template +struct NeighborhoodAttentionReferenceMask { + using NADim = NADim_; + using Causal = Causal_; + static_assert(rank(NADim{}) >= 1 && rank(NADim{}) < 4); + static_assert(rank(Causal{}) >= 1 && rank(Causal{}) < 4); + + NADim window_size; + NADim window_left; + NADim window_right; + NADim stride; + NADim dilation; + QKVLayout qkv_layout; + int num_additional_kv; + int additional_kv_offset; + + CUTE_DEVICE NeighborhoodAttentionReferenceMask( + NADim window_size_, + NADim stride_, + NADim dilation_, + QKVLayout qkv_layout_, + int num_additional_kv_) + : window_size(window_size_), + stride(stride_), + dilation(dilation_), + qkv_layout(qkv_layout_), + num_additional_kv(num_additional_kv_), + additional_kv_offset(size(qkv_layout_.shape())) { + window_left = get_window_left(window_size_); + window_right = get_window_right(window_size_); + } + + template + CUTE_DEVICE void apply_mask(AccQK& acc_qk, IndexQK const& index_qk) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(acc_qk); i++) { + // NOTE: this is very inefficient if there's multiple queries in the + // accumulator. But in the reference kernel we know accum size is 1x1. + auto [q_idx, kv_idx] = index_qk(i); + + // If kv_idx > size(qkv_shape), then it's part of the cross attention + // tokens. Rule is: allow all of them up to num_additional_kv, mask out + // the rest. + if (kv_idx >= additional_kv_offset && + kv_idx - additional_kv_offset < num_additional_kv) { + continue; + } else if (kv_idx >= additional_kv_offset) { + acc_qk(i) = -INFINITY; + continue; + } + + auto [q_coord, q_di_coord] = + map_idx_to_di_coords(q_idx, dilation, qkv_layout); + auto [kv_coord, kv_di_coord] = + map_idx_to_di_coords(kv_idx, dilation, qkv_layout); + + // Fixup input shape according to dilation group + auto qkv_shape = correct_qkv_shape_wrt_dilation( + qkv_layout.shape(), dilation, q_di_coord); + + auto kv_start = get_window_start( + q_coord, window_left, window_right, stride, qkv_shape); + auto kv_end = + get_window_end(q_coord, kv_start, window_size, qkv_shape); + + if (not is_equal(q_di_coord, kv_di_coord) or + not is_neighbor(kv_coord, kv_start, kv_end)) { + acc_qk(i) = -INFINITY; + } + } + } +}; + +} // namespace mask + +} // namespace reference +} // namespace cuda +} // namespace natten + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/reference/utils.hpp b/natten/csrc/include/natten/cuda/reference/utils.hpp new file mode 100644 index 00000000..b255ba6a --- /dev/null +++ b/natten/csrc/include/natten/cuda/reference/utils.hpp @@ -0,0 +1,65 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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 "cutlass/cutlass.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace natten { +namespace cuda { +namespace reference { + +// Make a stride compatible with torch's default contiguous layout. +template +CUTE_HOST_DEVICE auto make_qkv_stride(NADim qkv_shape) { + using namespace cute; + + static_assert(rank(qkv_shape) > 0 && rank(qkv_shape) < 4); + + if constexpr (rank(qkv_shape) == 1) { + return make_stride(_1{}); + } else if constexpr (rank(qkv_shape) == 2) { + return make_stride(get<1>(qkv_shape), _1{}); + } else { + return make_stride( + get<1>(qkv_shape) * get<2>(qkv_shape), get<2>(qkv_shape), _1{}); + } +} + +} // namespace reference +} // namespace cuda +} // namespace natten + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/cuda/tokperm/layouts.hpp b/natten/csrc/include/natten/cuda/tokperm/layouts.hpp new file mode 100644 index 00000000..eb3c78f6 --- /dev/null +++ b/natten/csrc/include/natten/cuda/tokperm/layouts.hpp @@ -0,0 +1,150 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ + +#pragma once + +#include +#include +#include +#include +#include + +#include + +#include +#include + +namespace natten::tokperm { + +using namespace cute; +using namespace natten::tokperm::utils; + +template +CUTE_HOST_DEVICE constexpr auto make_token_permuted_layout( + CuteTuple const& rest, + CuteTuple const& tile_shape, + CuteTuple const& dilation, + int batch, + int heads, + int dim, + bool flip_tiled_dims) { + static_assert(tuple_size_v >= 1 && tuple_size_v <= 3); + + if constexpr (tuple_size_v == 1) { + auto [R0] = rest; + auto [T0] = tile_shape; + auto [D0] = dilation; + + auto problem_shape_out = + cute::make_tuple(batch, cute::make_tuple(R0, T0, D0), heads, dim); + + // flip_tiled_dims doesn't make a difference for 1D + // (R0 T0 D0) -> (D0 R0 T0) + auto problem_shape_out_view = permute_tokens<2, 0, 1>(problem_shape_out); + + auto stride_out_view = + utils::make_torch_contiguous_stride(problem_shape_out_view); + + // (D0 R0 T0) -> (R0 T0 D0) + auto stride_out = permute_tokens<1, 2, 0>(stride_out_view); + + return make_layout(problem_shape_out, stride_out); + + } else if constexpr (tuple_size_v == 2) { + auto [R0, R1] = rest; + auto [T0, T1] = tile_shape; + auto [D0, D1] = dilation; + + auto problem_shape_out = cute::make_tuple( + batch, cute::make_tuple(R0, T0, D0, R1, T1, D1), heads, dim); + + if (flip_tiled_dims) { + // (R0 T0 D0 R1 T1 D1) -> (D1 D0 R1 R0 T1 T0) + auto problem_shape_out_view = + permute_tokens<5, 2, 3, 0, 4, 1>(problem_shape_out); + + auto stride_out_view = + utils::make_torch_contiguous_stride(problem_shape_out_view); + + // (D1 D0 R1 R0 T1 T0) -> (R0 T0 D0 R1 T1 D1) + auto stride_out = permute_tokens<3, 5, 1, 2, 4, 0>(stride_out_view); + + return make_layout(problem_shape_out, stride_out); + + } else { + // (R0 T0 D0 R1 T1 D1) -> (D0 D1 R0 R1 T0 T1) + auto problem_shape_out_view = + permute_tokens<2, 5, 0, 3, 1, 4>(problem_shape_out); + + auto stride_out_view = + utils::make_torch_contiguous_stride(problem_shape_out_view); + + // (D0 D1 R0 R1 T0 T1) -> (R0 T0 D0 R1 T1 D1) + auto stride_out = permute_tokens<2, 4, 0, 3, 5, 1>(stride_out_view); + + return make_layout(problem_shape_out, stride_out); + } + + } else if constexpr (tuple_size_v == 3) { + auto [R0, R1, R2] = rest; + auto [T0, T1, T2] = tile_shape; + auto [D0, D1, D2] = dilation; + + auto problem_shape_out = cute::make_tuple( + batch, + cute::make_tuple(R0, T0, D0, R1, T1, D1, R2, T2, D2), + heads, + dim); + + if (flip_tiled_dims) { + // (R0 T0 D0 R1 T1 D1 R2 T2 D2) -> (D2 D1 D0 R2 R1 R0 T2 T1 T0) + auto problem_shape_out_view = + permute_tokens<8, 5, 2, 6, 3, 0, 7, 4, 1>(problem_shape_out); + + auto stride_out_view = + utils::make_torch_contiguous_stride(problem_shape_out_view); + + // (D2 D1 D0 R2 R1 R0 T2 T1 T0) -> (R0 T0 D0 R1 T1 D1 R2 T2 D2) + auto stride_out = + permute_tokens<5, 8, 2, 4, 7, 1, 3, 6, 0>(stride_out_view); + + return make_layout(problem_shape_out, stride_out); + + } else { + // (R0 T0 D0 R1 T1 D1 R2 T2 D2) -> (D0 D1 D2 R0 R1 R2 T0 T1 T2) + auto problem_shape_out_view = + permute_tokens<2, 5, 8, 0, 3, 6, 1, 4, 7>(problem_shape_out); + + auto stride_out_view = + utils::make_torch_contiguous_stride(problem_shape_out_view); + + // (D0 D1 D2 R0 R1 R2 T0 T1 T2) -> (R0 T0 D0 R1 T1 D1 R2 T2 D2) + auto stride_out = + permute_tokens<3, 6, 0, 4, 7, 1, 5, 8, 2>(stride_out_view); + + return make_layout(problem_shape_out, stride_out); + } + } +} + +} // namespace natten::tokperm diff --git a/natten/csrc/include/natten/cuda/tokperm/token_permute_kernel.cuh b/natten/csrc/include/natten/cuda/tokperm/token_permute_kernel.cuh new file mode 100644 index 00000000..5bff3e02 --- /dev/null +++ b/natten/csrc/include/natten/cuda/tokperm/token_permute_kernel.cuh @@ -0,0 +1,355 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ + +#pragma once + +#include +#include +#include +#include + +#include +#include +#include + +namespace natten::tokperm::kernel { + +using namespace cute; + +template +CUTE_HOST_DEVICE constexpr auto perm2unperm( + DstCrd const& dstCrd, + Dim const& rest, + Dim const& tile, + Dim const& dilation) { + static_assert( + rank(DstCrd{}) == 3 || rank(DstCrd{}) == 6 || rank(DstCrd{}) == 9); + static_assert(rank(DstCrd{}) / 3 == rank(Dim{})); + + if constexpr (rank(Dim{}) == 1) { + auto tiled_layout_0 = make_layout( + make_shape(get<0>(rest), get<0>(tile), get<0>(dilation)), + make_stride(get<0>(tile) * get<0>(dilation), get<0>(dilation), _1{})); + + auto crd_0 = select<0, 1, 2>(dstCrd); + auto crd_src_0 = tiled_layout_0(crd_0); + return make_tuple(crd_src_0); + } else if constexpr (rank(Dim{}) == 2) { + auto tiled_layout_0 = make_layout( + make_shape(get<0>(rest), get<0>(tile), get<0>(dilation)), + make_stride(get<0>(tile) * get<0>(dilation), get<0>(dilation), _1{})); + auto tiled_layout_1 = make_layout( + make_shape(get<1>(rest), get<1>(tile), get<1>(dilation)), + make_stride(get<1>(tile) * get<1>(dilation), get<1>(dilation), _1{})); + + auto crd_0 = select<0, 1, 2>(dstCrd); + auto crd_1 = select<3, 4, 5>(dstCrd); + auto crd_src_0 = tiled_layout_0(crd_0); + auto crd_src_1 = tiled_layout_1(crd_1); + return make_tuple(crd_src_0, crd_src_1); + } else if constexpr (rank(Dim{}) == 3) { + auto tiled_layout_0 = make_layout( + make_shape(get<0>(rest), get<0>(tile), get<0>(dilation)), + make_stride(get<0>(tile) * get<0>(dilation), get<0>(dilation), _1{})); + auto tiled_layout_1 = make_layout( + make_shape(get<1>(rest), get<1>(tile), get<1>(dilation)), + make_stride(get<1>(tile) * get<1>(dilation), get<1>(dilation), _1{})); + auto tiled_layout_2 = make_layout( + make_shape(get<2>(rest), get<2>(tile), get<2>(dilation)), + make_stride(get<2>(tile) * get<2>(dilation), get<2>(dilation), _1{})); + + auto crd_0 = select<0, 1, 2>(dstCrd); + auto crd_1 = select<3, 4, 5>(dstCrd); + auto crd_2 = select<6, 7, 8>(dstCrd); + auto crd_src_0 = tiled_layout_0(crd_0); + auto crd_src_1 = tiled_layout_1(crd_1); + auto crd_src_2 = tiled_layout_2(crd_2); + return make_tuple(crd_src_0, crd_src_1, crd_src_2); + } +} + +template +CUTE_HOST_DEVICE constexpr auto unperm2perm( + DstCrd const& dstCrd, + Dim const& rest, + Dim const& tile, + Dim const& dilation) { + static_assert(rank(DstCrd{}) == rank(Dim{})); + + if constexpr (rank(Dim{}) == 1) { + auto tiled_layout_0 = make_layout( + make_shape(get<0>(rest), get<0>(tile), get<0>(dilation)), + make_stride(get<0>(tile) * get<0>(dilation), get<0>(dilation), _1{})); + + auto [crd_0] = dstCrd; + auto [crd_src_0_r, crd_src_0_t, crd_src_0_d] = + idx2crd(crd_0, tiled_layout_0.shape(), tiled_layout_0.stride()); + return make_tuple(crd_src_0_r, crd_src_0_t, crd_src_0_d); + } else if constexpr (rank(Dim{}) == 2) { + auto tiled_layout_0 = make_layout( + make_shape(get<0>(rest), get<0>(tile), get<0>(dilation)), + make_stride(get<0>(tile) * get<0>(dilation), get<0>(dilation), _1{})); + auto tiled_layout_1 = make_layout( + make_shape(get<1>(rest), get<1>(tile), get<1>(dilation)), + make_stride(get<1>(tile) * get<1>(dilation), get<1>(dilation), _1{})); + + auto [crd_0, crd_1] = dstCrd; + auto [crd_src_0_r, crd_src_0_t, crd_src_0_d] = + idx2crd(crd_0, tiled_layout_0.shape(), tiled_layout_0.stride()); + auto [crd_src_1_r, crd_src_1_t, crd_src_1_d] = + idx2crd(crd_1, tiled_layout_1.shape(), tiled_layout_1.stride()); + return make_tuple( + crd_src_0_r, + crd_src_0_t, + crd_src_0_d, + crd_src_1_r, + crd_src_1_t, + crd_src_1_d); + } else if constexpr (rank(Dim{}) == 3) { + auto tiled_layout_0 = make_layout( + make_shape(get<0>(rest), get<0>(tile), get<0>(dilation)), + make_stride(get<0>(tile) * get<0>(dilation), get<0>(dilation), _1{})); + auto tiled_layout_1 = make_layout( + make_shape(get<1>(rest), get<1>(tile), get<1>(dilation)), + make_stride(get<1>(tile) * get<1>(dilation), get<1>(dilation), _1{})); + auto tiled_layout_2 = make_layout( + make_shape(get<2>(rest), get<2>(tile), get<2>(dilation)), + make_stride(get<2>(tile) * get<2>(dilation), get<2>(dilation), _1{})); + + auto [crd_0, crd_1, crd_2] = dstCrd; + auto [crd_src_0_r, crd_src_0_t, crd_src_0_d] = + idx2crd(crd_0, tiled_layout_0.shape(), tiled_layout_0.stride()); + auto [crd_src_1_r, crd_src_1_t, crd_src_1_d] = + idx2crd(crd_1, tiled_layout_1.shape(), tiled_layout_1.stride()); + auto [crd_src_2_r, crd_src_2_t, crd_src_2_d] = + idx2crd(crd_2, tiled_layout_2.shape(), tiled_layout_2.stride()); + return make_tuple( + crd_src_0_r, + crd_src_0_t, + crd_src_0_d, + crd_src_1_r, + crd_src_1_t, + crd_src_1_d, + crd_src_2_r, + crd_src_2_t, + crd_src_2_d); + } +} + +template < + typename TokenShapeIn, + typename TokenShapeOut, + class ElementIn, + class ElementOut, + bool IsUnpermute = false, + int kElementsPerLoad = 4> +struct TokenPermuteKernel { + using OffsetTypeInternal = uint64_t; + using TokenShape = + cute::conditional_t; + + // B, token layout shape, tile shape, dilation, H, D + using ProblemShapeIn = cute::tuple; + using ProblemShapeOut = cute::tuple; + + using TokenStrideIn = utils::make_tuple_type; + using TokenStrideOut = utils::make_tuple_type; + using StrideIn = cute::tuple; + using StrideOut = cute::tuple; + + using ClusterShape = Shape<_1, _1, _1>; + static constexpr int SharedStorageSize = 0; + + struct Arguments { + ProblemShapeIn problem_shape_src; + ProblemShapeOut problem_shape_dst; + + ElementIn* ptr_src; + ElementOut* ptr_dst; + + StrideIn stride_src; + StrideOut stride_dst; + + TokenShape rest; + TokenShape tile; + TokenShape dilation; + }; + + using Params = Arguments; + + static const int MinBlocksPerMultiprocessor = 1; + static const int MaxThreadsPerBlock = 128; + + static const int kBlockSeq = 8; + + static const int kNumThreadsD = 16; + static const int kNumThreadsSeq = MaxThreadsPerBlock / kNumThreadsD; + + static const int kIterationsSeq = kBlockSeq / kNumThreadsSeq; + + static bool can_implement(Arguments const& args) { + if (cute::size<0>(args.problem_shape_src) != + cute::size<0>(args.problem_shape_dst)) { + std::cerr + << "Token Permute/UnPermute requires input and output batch sizes to match." + << std::endl; + return false; + } + if (IsUnpermute) { + if (cute::size<1>(args.problem_shape_src) < + cute::size<1>(args.problem_shape_dst)) { + std::cerr << "Token UnPermute requires more or equal input tokens." + << std::endl; + return false; + } + } else { + if (cute::size<1>(args.problem_shape_src) > + cute::size<1>(args.problem_shape_dst)) { + std::cerr + << "Token Permute/UnPermute requires more or equal output tokens." + << std::endl; + return false; + } + } + if (cute::size<2>(args.problem_shape_src) != + cute::size<2>(args.problem_shape_dst)) { + std::cerr + << "Token Permute/UnPermute requires input and output heads to match." + << std::endl; + return false; + } + if (cute::size<3>(args.problem_shape_src) != + cute::size<3>(args.problem_shape_dst)) { + std::cerr + << "Token Permute/UnPermute requires input and output head dims to match." + << std::endl; + return false; + } + if (cute::get<3>(args.stride_src) != _1{} || + cute::get<3>(args.stride_dst) != _1{}) { + std::cerr + << "Token Permute/UnPermute requires contiguous input and output (stride[-1] == 1)." + << std::endl; + return false; + } + if (cute::size<3>(args.problem_shape_src) % kElementsPerLoad != 0) { + std::cerr + << "Token Permute/UnPermute requires head dim to be evenly divisible by load size (" + << kElementsPerLoad << ")" << std::endl; + return false; + } + + return true; + } + + static dim3 get_grid_shape(Params const& params) { + dim3 grid( + // never put seq in z, long seqs can easily exceed the 64K limit + cute::ceil_div(size<1>(params.problem_shape_dst), kBlockSeq), + cute::size<2>(params.problem_shape_dst), + cute::size<0>(params.problem_shape_dst)); + return grid; + } + + static dim3 get_block_shape() { + dim3 block(kNumThreadsSeq, kNumThreadsD, 1); + return block; + } + + static Params to_underlying_arguments(Arguments const& args) { + return args; + } + + CUTLASS_DEVICE void operator()(const Params& params, char* smem) { + auto ptr_src_bh = params.ptr_src + + (static_cast(get<0>(params.stride_src)) * + static_cast(blockIdx.z) + + static_cast(get<2>(params.stride_src)) * + static_cast(blockIdx.y)); + auto ptr_dst_bh = params.ptr_dst + + (static_cast(get<0>(params.stride_dst)) * + static_cast(blockIdx.z) + + static_cast(get<2>(params.stride_dst)) * + static_cast(blockIdx.y)); + + auto token_layout_src = make_layout( + get<1>(params.problem_shape_src), get<1>(params.stride_src)); + auto token_layout_dst = make_layout( + get<1>(params.problem_shape_dst), get<1>(params.stride_dst)); + + auto src_shape = token_layout_src.shape(); + + auto seqlen = size<1>(params.problem_shape_dst); + + for (int idx_s_t = threadIdx.x; idx_s_t < kBlockSeq; + idx_s_t += kNumThreadsSeq) { + int idx_s = idx_s_t + kBlockSeq * blockIdx.x; + if (idx_s >= seqlen) + continue; + + auto crd_dst = idx2crd(idx_s, token_layout_dst.shape()); + bool pred = false; + + TokenShapeIn crd_src; + if constexpr (IsUnpermute) { + crd_src = + unperm2perm(crd_dst, params.rest, params.tile, params.dilation); + } else { + crd_src = + perm2unperm(crd_dst, params.rest, params.tile, params.dilation); + } + + pred = elem_less(crd_src, src_shape); + + auto ptr_src_bhs = ptr_src_bh + token_layout_src(crd_src); + auto ptr_dst_bhs = ptr_dst_bh + token_layout_dst(crd_dst); + + for (int idx_d = threadIdx.y * kElementsPerLoad; + idx_d < get<3>(params.problem_shape_dst); + idx_d += kElementsPerLoad * kNumThreadsD) { + ElementIn value_src[kElementsPerLoad]; + ElementOut value_dst[kElementsPerLoad]; + + using VecSrc = uint_bit_t * kElementsPerLoad>; + using VecDst = uint_bit_t * kElementsPerLoad>; + if (pred) { + *reinterpret_cast(value_src) = + *reinterpret_cast(&ptr_src_bhs[idx_d]); + + for (int v = 0; v < kElementsPerLoad; v++) { + value_dst[v] = static_cast(value_src[v]); + } + } else { + for (int v = 0; v < kElementsPerLoad; v++) { + value_dst[v] = static_cast(0); + } + } + + *reinterpret_cast(&ptr_dst_bhs[idx_d]) = + *reinterpret_cast(value_dst); + } + } + } +}; + +} // namespace natten::tokperm::kernel diff --git a/natten/csrc/include/natten/cuda/tokperm/tokperm.hpp b/natten/csrc/include/natten/cuda/tokperm/tokperm.hpp new file mode 100644 index 00000000..6e6c14c9 --- /dev/null +++ b/natten/csrc/include/natten/cuda/tokperm/tokperm.hpp @@ -0,0 +1,262 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ + +#pragma once + +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include + +namespace natten::tokperm { + +using namespace cute; + +template +bool token_permute_op( + ElementIn* ptr_in, + ElementOut* ptr_out, + int batch, + int heads, + int seqlen_out, + int dim, + CuteTuple& token_layout, + CuteTuple& tile_shape, + CuteTuple& dilation, + bool flip_tiled_dims, + cudaStream_t stream) { + static constexpr int NumDims = tuple_size_v; + + // NumDims = 1: tuple + // NumDims = 2: tuple + // NumDims = 3: tuple + using DimIn = utils::make_tuple_type; + // (token mode/dim) -> (tile, dilation, rest) + using DimOut = utils::make_tuple_type; + + auto rest = ceil_div(ceil_div(token_layout, tile_shape), dilation); + + auto problem_shape_in = cute::make_tuple(batch, token_layout, heads, dim); + auto stride_in = utils::make_torch_contiguous_stride(problem_shape_in); + + auto layout_out = make_token_permuted_layout( + rest, tile_shape, dilation, batch, heads, dim, flip_tiled_dims); + + auto problem_shape_out = layout_out.shape(); + auto stride_out = layout_out.stride(); + + if (cute::size<1>(problem_shape_out) != seqlen_out) { + std::cerr << "Token Permute output must have sequence length of exactly " + << cute::size<1>(problem_shape_out) + << ", got tensor with sequence length " << seqlen_out << "!" + << std::endl; + return false; + } + + using OperationAlign1 = kernel::TokenPermuteKernel< + DimIn, + DimOut, + ElementIn, + ElementOut, + /* IsUnpermute = */ false, + 1>; + using OperationAlign4 = kernel::TokenPermuteKernel< + DimIn, + DimOut, + ElementIn, + ElementOut, + /* IsUnpermute = */ false, + 4>; + using OperationAlign8 = kernel::TokenPermuteKernel< + DimIn, + DimOut, + ElementIn, + ElementOut, + /* IsUnpermute = */ false, + 8>; + + auto launch_kernel = [&](auto& op) { + using Operation = std::remove_reference_t; + using Arguments = typename Operation::Arguments; + + Arguments arguments{ + problem_shape_in, + problem_shape_out, + ptr_in, + ptr_out, + stride_in, + stride_out, + rest, + tile_shape, + dilation, + }; + + if (not op.can_implement(arguments)) { + std::cerr << "Token Permute kernel is not supported." << std::endl; + return false; + } + + auto params = Operation::to_underlying_arguments(arguments); + auto grid = Operation::get_grid_shape(params); + auto block = Operation::get_block_shape(); + int smem_size = Operation::SharedStorageSize; + cutlass::device_kernel + <<>>(params); + + cudaError_t result = cudaGetLastError(); + if (result != cudaSuccess) { + std::cerr << "Failed to launch Token Permute kernel. Last CUDA error is: " + << cudaGetErrorString(result) << std::endl; + return false; + } + + return true; + }; + + if (dim % 8 == 0) { + OperationAlign8 op; + return launch_kernel(op); + } else if (dim % 4 == 0) { + OperationAlign4 op; + return launch_kernel(op); + } else { + OperationAlign1 op; + return launch_kernel(op); + } +} + +template +bool token_unpermute_op( + ElementIn* ptr_in, + ElementOut* ptr_out, + int batch, + int heads, + int dim, + CuteTuple& token_layout, + CuteTuple& tile_shape, + CuteTuple& dilation, + bool flip_tiled_dims, + cudaStream_t stream) { + static constexpr int NumDims = tuple_size_v; + + // NumDims = 1: tuple + // NumDims = 2: tuple + // NumDims = 3: tuple + using DimOut = utils::make_tuple_type; + // (token mode/dim) -> (tile, dilation, rest) + using DimIn = utils::make_tuple_type; + + auto rest = ceil_div(ceil_div(token_layout, tile_shape), dilation); + + auto problem_shape_out = cute::make_tuple(batch, token_layout, heads, dim); + auto stride_out = utils::make_torch_contiguous_stride(problem_shape_out); + + auto layout_in = make_token_permuted_layout( + rest, tile_shape, dilation, batch, heads, dim, flip_tiled_dims); + + auto problem_shape_in = layout_in.shape(); + auto stride_in = layout_in.stride(); + + using OperationAlign1 = kernel::TokenPermuteKernel< + DimIn, + DimOut, + ElementIn, + ElementOut, + /* IsUnpermute = */ true, + 1>; + using OperationAlign4 = kernel::TokenPermuteKernel< + DimIn, + DimOut, + ElementIn, + ElementOut, + /* IsUnpermute = */ true, + 4>; + using OperationAlign8 = kernel::TokenPermuteKernel< + DimIn, + DimOut, + ElementIn, + ElementOut, + /* IsUnpermute = */ true, + 8>; + + auto launch_kernel = [&](auto& op) { + using Operation = std::remove_reference_t; + using Arguments = typename Operation::Arguments; + + Arguments arguments{ + problem_shape_in, + problem_shape_out, + ptr_in, + ptr_out, + stride_in, + stride_out, + rest, + tile_shape, + dilation, + }; + + if (not op.can_implement(arguments)) { + std::cerr << "Token UnPermute kernel is not supported." << std::endl; + return false; + } + + auto params = Operation::to_underlying_arguments(arguments); + auto grid = Operation::get_grid_shape(params); + auto block = Operation::get_block_shape(); + int smem_size = Operation::SharedStorageSize; + cutlass::device_kernel + <<>>(params); + + cudaError_t result = cudaGetLastError(); + if (result != cudaSuccess) { + std::cerr + << "Failed to launch Token UnPermute kernel. Last CUDA error is: " + << cudaGetErrorString(result) << std::endl; + return false; + } + + return true; + }; + + if (dim % 8 == 0) { + OperationAlign8 op; + return launch_kernel(op); + } else if (dim % 4 == 0) { + OperationAlign4 op; + return launch_kernel(op); + } else { + OperationAlign1 op; + return launch_kernel(op); + } +} + +} // namespace natten::tokperm diff --git a/natten/csrc/include/natten/cuda/tokperm/utils/permute.cuh b/natten/csrc/include/natten/cuda/tokperm/utils/permute.cuh new file mode 100644 index 00000000..fbda3e35 --- /dev/null +++ b/natten/csrc/include/natten/cuda/tokperm/utils/permute.cuh @@ -0,0 +1,57 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ + +#pragma once + +#include +#include +#include +#include +#include + +#include + +#include + +namespace natten::tokperm::utils { + +using namespace cute; + +// `permute_tokens`: utility for applying permutations to arbitrary-sized token +// views. Expects rank-4 inputs (batch, (*tokens), heads, head_dim) +// +// NOTE: does not check Permutation to ensure it's a valid permutation +// I.e. permute_tokens<1,1> repeats token dimension 1 without raising an error. +template +CUTE_HOST_DEVICE constexpr ProblemShape permute_tokens( + ProblemShape problem_shape) { + static_assert(rank(ProblemShape{}) == 4); + + return cute::make_tuple( + get<0>(problem_shape), + select(get<1>(problem_shape)), + get<2>(problem_shape), + get<3>(problem_shape)); +} + +} // namespace natten::tokperm::utils diff --git a/natten/csrc/include/natten/cuda/tokperm/utils/stride.cuh b/natten/csrc/include/natten/cuda/tokperm/utils/stride.cuh new file mode 100644 index 00000000..7388e2d7 --- /dev/null +++ b/natten/csrc/include/natten/cuda/tokperm/utils/stride.cuh @@ -0,0 +1,75 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ + +#pragma once + +#include +#include + +#include + +namespace natten::tokperm::utils { + +using namespace cute; + +namespace detail { + +template +struct StrideHelper { + static_assert(cute::rank(Shape{}) >= 2); + static_assert(depth(Shape{}) == 1); + static constexpr int Rank = cute::rank(Shape{}); + + using StrideType_ = utils::make_tuple_type; + // stride for heads can be int32 + using StrideType = decltype(append(append(StrideType_{}, int{}), _1{})); + + CUTE_HOST_DEVICE static constexpr StrideType make_stride(Shape const& shape) { + StrideType stride; + + get(stride) = _1{}; // dim + get(stride) = get(shape); // heads + cute::for_each(cute::make_range<0, Rank - 1>{}, [&](auto i) { + static_assert(i < Rank - 1); + get(stride) = + static_cast(get(shape)) * + get(stride); + }); + + return stride; + } +}; + +} // namespace detail + +template +CUTE_HOST_DEVICE static constexpr auto make_torch_contiguous_stride( + Shape const& shape) { + auto flattened_shape = cute::flatten(shape); + auto flattened_stride = + detail::StrideHelper::make_stride( + flattened_shape); + return cute::unflatten(flattened_stride, shape); +} + +} // namespace natten::tokperm::utils diff --git a/natten/csrc/include/natten/cuda/tokperm/utils/tuple.cuh b/natten/csrc/include/natten/cuda/tokperm/utils/tuple.cuh new file mode 100644 index 00000000..0353119e --- /dev/null +++ b/natten/csrc/include/natten/cuda/tokperm/utils/tuple.cuh @@ -0,0 +1,40 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ + +#pragma once + +#include +#include +#include + +namespace natten::tokperm::utils { + +template +auto make_tuple_type_impl(cute::index_sequence) + -> cute::tuple; + +template +using make_tuple_type = + decltype(make_tuple_type_impl(cute::make_index_sequence{})); + +} // namespace natten::tokperm::utils diff --git a/natten/csrc/include/natten/cuda/utils/cuda.h b/natten/csrc/include/natten/cuda/utils/cuda.h new file mode 100644 index 00000000..dad7a9a6 --- /dev/null +++ b/natten/csrc/include/natten/cuda/utils/cuda.h @@ -0,0 +1,37 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ + +#pragma once + +#include +#include + +#define NATTEN_CUDA_CHECK(status) \ + [&] { \ + cudaError_t error = status; \ + if (error != cudaSuccess) { \ + throw std::runtime_error( \ + std::string("NATTEN failure: CUDA runtime error: ") + \ + cudaGetErrorString(error) + " at: " + std::to_string(__LINE__)); \ + } \ + }(); diff --git a/natten/csrc/include/natten/cuda/utils/cutlass.cuh b/natten/csrc/include/natten/cuda/utils/cutlass.cuh new file mode 100644 index 00000000..446d03a3 --- /dev/null +++ b/natten/csrc/include/natten/cuda/utils/cutlass.cuh @@ -0,0 +1,112 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + **************************************************************************************************/ +/*************************************************************************************************** + * Copyright (c) 2017 - 2026 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 +#include + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { + +/// SM90 Device Kernel wrapper +template +CUTLASS_GLOBAL +#ifdef __CUDACC__ +// Enclosing this in __CUDACC__ suppresses MSVC warnings. +__launch_bounds__( + Operator::MaxThreadsPerBlock, + Operator::MinBlocksPerMultiprocessor) +#endif // __CUDACC__ + void device_kernel_sm90(CUTLASS_GRID_CONSTANT + typename Operator::Params const params) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ == 900 + // Dynamic shared memory base pointer + extern __shared__ char smem[]; + Operator op; + op(params, smem); + cutlass::arch::synclog_print(); +#else + printf( + "FATAL: This kernel was built for SM90, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +#endif +} + +/// SM100 Device Kernel wrapper +template +CUTLASS_GLOBAL +#ifdef __CUDACC__ +// Enclosing this in __CUDACC__ suppresses MSVC warnings. +__launch_bounds__( + Operator::MaxThreadsPerBlock, + Operator::MinBlocksPerMultiprocessor) +#endif // __CUDACC__ + void device_kernel_sm100(CUTLASS_GRID_CONSTANT + typename Operator::Params const params) { +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1030 + // Dynamic shared memory base pointer + extern __shared__ char smem[]; + Operator op; + op(params, smem); + cutlass::arch::synclog_print(); +#else + printf( + "FATAL: This kernel was built for SM100 and SM103, but attempted to launch from SM%d\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +#endif +} + +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// + +#define NATTEN_CUTLASS_CHECK(status) \ + [&] { \ + cutlass::Status error = status; \ + if (error != cutlass::Status::kSuccess) { \ + throw std::runtime_error( \ + std::string("NATTEN failure: cutlass error: ") + \ + cutlassGetStatusString(error) + " at: " + std::to_string(__LINE__)); \ + } \ + }(); diff --git a/natten/csrc/include/natten/cuda/utils/generic_cutlass_device.hpp b/natten/csrc/include/natten/cuda/utils/generic_cutlass_device.hpp new file mode 100644 index 00000000..99112302 --- /dev/null +++ b/natten/csrc/include/natten/cuda/utils/generic_cutlass_device.hpp @@ -0,0 +1,288 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 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 "cute/tensor.hpp" +#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::device { + +//////////////////////////////////////////////////////////////////////////////// +////////////////////////////// CUTLASS 3.x API ///////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +template +class DeviceKernel { + 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("DeviceKernel::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( + "DeviceKernel::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("DeviceKernel()::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) { + CUTLASS_TRACE_HOST("DeviceKernel::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); + } else { + launch_result = Status::kSuccess; + cutlass::arch::synclog_setup(); + 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::device + +//////////////////////////////////////////////////////////////////////////////// diff --git a/natten/csrc/include/natten/fmha.h b/natten/csrc/include/natten/fmha.h new file mode 100644 index 00000000..0829fd13 --- /dev/null +++ b/natten/csrc/include/natten/fmha.h @@ -0,0 +1,74 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ +/*! \file + \brief FMHA interface +*/ + +#pragma once +#include + +#include + +namespace natten { + +void fmha_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::optional& logsumexp, + bool is_causal, + float attn_scale, + int32_t query_tile_size, + int32_t key_tile_size, + // varlen + const at::optional& cumulative_seqlen_Q, + const at::optional& cumulative_seqlen_KV, + // only used if cumulative_seqlen_Q and cumulative_seqlen_KV are specified + int max_seqlen_Q, + int max_seqlen_KV); + +void fmha_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + bool is_causal, + float attn_scale, + int32_t query_tile_size, + int32_t key_tile_size, + int32_t num_splits_key, + bool compute_delta_with_torch, + // varlen + const at::optional& cumulative_seqlen_Q, + const at::optional& cumulative_seqlen_KV, + // only used if cumulative_seqlen_Q and cumulative_seqlen_KV are specified + int max_seqlen_Q, + int max_seqlen_KV); + +} // namespace natten diff --git a/natten/csrc/include/natten/fna.h b/natten/csrc/include/natten/fna.h new file mode 100644 index 00000000..3fb495d1 --- /dev/null +++ b/natten/csrc/include/natten/fna.h @@ -0,0 +1,140 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ +/*! \file + \brief FNA interface +*/ + +#pragma once +#include + +#include + +namespace natten { + +// Forward + +void na1d_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::optional& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& query_tile_size, + const std::tuple& key_tile_size); + +void na2d_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::optional& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& query_tile_size, + const std::tuple& key_tile_size); + +void na3d_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::optional& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& query_tile_size, + const std::tuple& key_tile_size); + +// Backward + +void na1d_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& query_tile_size, + const std::tuple& key_tile_size, + const std::tuple& num_splits_key, + bool compute_delta_with_torch); + +void na2d_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& query_tile_size, + const std::tuple& key_tile_size, + const std::tuple& num_splits_key, + bool compute_delta_with_torch); + +void na3d_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& query_tile_size, + const std::tuple& key_tile_size, + const std::tuple& num_splits_key, + bool compute_delta_with_torch); + +} // namespace natten diff --git a/natten/csrc/include/natten/helpers.h b/natten/csrc/include/natten/helpers.h new file mode 100644 index 00000000..05521770 --- /dev/null +++ b/natten/csrc/include/natten/helpers.h @@ -0,0 +1,461 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ + +#pragma once +#include + +#define CHECK_CONTIGUOUS(x) \ + TORCH_CHECK(not x.is_sparse(), #x " must be a dense tensor"); \ + TORCH_CHECK(x.is_contiguous(), #x " must be contiguous"); + +#define CHECK_CUDA(x) \ + TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor."); + +namespace natten { + +inline void AssertOddKernelSize(int32_t kernel_size) { + TORCH_CHECK( + kernel_size % 2 == 1, + "This operation only supports odd-sized kernel sizes, got ", + kernel_size, + "."); +} + +inline void AssertOddKernelSize(const std::tuple& kernel_size) { + AssertOddKernelSize(std::get<0>(kernel_size)); +} + +inline void AssertOddKernelSize( + const std::tuple& kernel_size) { + AssertOddKernelSize(std::get<0>(kernel_size)); + AssertOddKernelSize(std::get<1>(kernel_size)); +} + +inline void AssertOddKernelSize( + const std::tuple& kernel_size) { + AssertOddKernelSize(std::get<0>(kernel_size)); + AssertOddKernelSize(std::get<1>(kernel_size)); + AssertOddKernelSize(std::get<2>(kernel_size)); +} + +inline void FnaRPBChecks(int32_t kernel_size, int32_t stride, bool has_rpb) { + TORCH_CHECK( + !has_rpb || (kernel_size % 2 == 1 && stride == 1), + "This operation only supports odd-sized kernel sizes and stride=1 with RPB enabled, got ", + kernel_size, + "."); +} + +inline void FnaRPBChecks( + const std::tuple& kernel_size, + const std::tuple& stride, + bool has_rpb) { + FnaRPBChecks(std::get<0>(kernel_size), std::get<0>(stride), has_rpb); +} + +inline void FnaRPBChecks( + const std::tuple& kernel_size, + const std::tuple& stride, + bool has_rpb) { + FnaRPBChecks(std::get<0>(kernel_size), std::get<0>(stride), has_rpb); + FnaRPBChecks(std::get<1>(kernel_size), std::get<1>(stride), has_rpb); +} + +inline void FnaRPBChecks( + const std::tuple& kernel_size, + const std::tuple& stride, + bool has_rpb) { + FnaRPBChecks(std::get<0>(kernel_size), std::get<0>(stride), has_rpb); + FnaRPBChecks(std::get<1>(kernel_size), std::get<1>(stride), has_rpb); + FnaRPBChecks(std::get<2>(kernel_size), std::get<2>(stride), has_rpb); +} + +inline void CheckArgs(int32_t kernel_size, int32_t dilation) { + TORCH_CHECK( + kernel_size > 1, + "Kernel size must be greater than 1, got ", + kernel_size, + "."); + TORCH_CHECK( + dilation >= 1, + "Dilation must be a positive integer, got ", + dilation, + "."); +} + +inline void CheckArgs( + const std::tuple& kernel_size, + const std::tuple& dilation) { + CheckArgs(std::get<0>(kernel_size), std::get<0>(dilation)); +} + +inline void CheckArgs( + const std::tuple& kernel_size, + const std::tuple& dilation) { + CheckArgs(std::get<0>(kernel_size), std::get<0>(dilation)); + CheckArgs(std::get<1>(kernel_size), std::get<1>(dilation)); +} + +inline void CheckArgs( + const std::tuple& kernel_size, + const std::tuple& dilation) { + CheckArgs(std::get<0>(kernel_size), std::get<0>(dilation)); + CheckArgs(std::get<1>(kernel_size), std::get<1>(dilation)); + CheckArgs(std::get<2>(kernel_size), std::get<2>(dilation)); +} + +inline void CheckArgs(int32_t kernel_size, int32_t stride, int32_t dilation) { + CheckArgs(kernel_size, dilation); + TORCH_CHECK( + stride >= 1, "Stride must be a positive integer, got ", stride, "."); + TORCH_CHECK( + stride <= kernel_size, + "Stride must be smaller than or equal to kernel size, got ", + "kernel_size=", + kernel_size, + ", stride=", + stride, + "."); +} + +inline void CheckArgs( + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation) { + CheckArgs( + std::get<0>(kernel_size), std::get<0>(stride), std::get<0>(dilation)); +} + +inline void CheckArgs( + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation) { + CheckArgs( + std::get<0>(kernel_size), std::get<0>(stride), std::get<0>(dilation)); + CheckArgs( + std::get<1>(kernel_size), std::get<1>(stride), std::get<1>(dilation)); +} + +inline void CheckArgs( + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation) { + CheckArgs( + std::get<0>(kernel_size), std::get<0>(stride), std::get<0>(dilation)); + CheckArgs( + std::get<1>(kernel_size), std::get<1>(stride), std::get<1>(dilation)); + CheckArgs( + std::get<2>(kernel_size), std::get<2>(stride), std::get<2>(dilation)); +} + +inline void CheckArgsAgainstDim( + int32_t dim, + int32_t kernel_size, + int32_t dilation) { + TORCH_CHECK( + kernel_size * dilation <= dim, + "Input axes must be greater than or equal to the product of kernel size and dilation. " + "Got kernel size ", + kernel_size, + ", dilation ", + dilation, + ", but dimension size was ", + dim, + "."); +} + +inline void CheckArgsAgainstDim( + const std::tuple& dim, + const std::tuple& kernel_size, + const std::tuple& dilation) { + CheckArgsAgainstDim( + std::get<0>(dim), std::get<0>(kernel_size), std::get<0>(dilation)); +} + +inline void CheckArgsAgainstDim( + const std::tuple& dim, + const std::tuple& kernel_size, + const std::tuple& dilation) { + CheckArgsAgainstDim( + std::get<0>(dim), std::get<0>(kernel_size), std::get<0>(dilation)); + CheckArgsAgainstDim( + std::get<1>(dim), std::get<1>(kernel_size), std::get<1>(dilation)); +} + +inline void CheckArgsAgainstDim( + const std::tuple& dim, + const std::tuple& kernel_size, + const std::tuple& dilation) { + CheckArgsAgainstDim( + std::get<0>(dim), std::get<0>(kernel_size), std::get<0>(dilation)); + CheckArgsAgainstDim( + std::get<1>(dim), std::get<1>(kernel_size), std::get<1>(dilation)); + CheckArgsAgainstDim( + std::get<2>(dim), std::get<2>(kernel_size), std::get<2>(dilation)); +} + +inline void CheckIfPropertiesMatch(const at::Tensor& a, const at::Tensor& b) { + TORCH_CHECK( + a.device().is_cuda() == b.device().is_cuda(), + "Expected all tensors to be on the same device."); + TORCH_CHECK( + a.scalar_type() == b.scalar_type(), "Input tensors must match in dtype!"); +} + +inline void CheckIfPropertiesMatch( + const at::Tensor& a, + const at::Tensor& b, + const at::Tensor& c) { + TORCH_CHECK( + a.device().is_cuda() == b.device().is_cuda() && + b.device().is_cuda() == c.device().is_cuda(), + "Expected all tensors to be on the same device."); + TORCH_CHECK( + a.scalar_type() == b.scalar_type() && b.scalar_type() == c.scalar_type(), + "Input tensors must match in dtype!"); +} + +template +void CheckIfTensorShapesMatch(const at::Tensor& a, const at::Tensor& b) { + static_assert(NaDim >= 1 && NaDim < 4); + static constexpr size_t Rank = NaDim + 3; + TORCH_CHECK( + a.dim() == b.dim() && a.dim() == Rank, "Expected ", Rank, "-D tensors."); + for (size_t i = 0; i < Rank; ++i) { + TORCH_CHECK( + a.size(i) == b.size(i), + "Tensor shape mismatch at dimension ", + i, + ": ", + a.size(i), + " != ", + b.size(i)); + } +} + +template +void CheckIfTensorShapesMatchExceptHeadDim( + const at::Tensor& a, + const at::Tensor& b) { + static_assert(NaDim >= 1 && NaDim < 4); + static constexpr size_t Rank = NaDim + 3; + TORCH_CHECK( + a.dim() == b.dim() && a.dim() == Rank, "Expected ", Rank, "-D tensors."); + for (size_t i = 0; i < Rank - 1; ++i) { + TORCH_CHECK( + a.size(i) == b.size(i), + "Tensor shape mismatch at dimension ", + i, + ": ", + a.size(i), + " != ", + b.size(i)); + } +} + +inline void CheckIfBatchHeadsMatch(const at::Tensor& a, const at::Tensor& b) { + TORCH_CHECK(a.dim() == b.dim(), "Expected tensors to match in rank."); + auto tensor_rank = a.dim(); + TORCH_CHECK( + a.size(0) == b.size(0), + "Tensors don't match in batch size; ", + a.size(0), + " != ", + b.size(0)); + TORCH_CHECK( + a.size(tensor_rank - 2) == b.size(tensor_rank - 2), + "Tensors don't match in number of heads; ", + a.size(tensor_rank - 2), + " != ", + b.size(tensor_rank - 2)); +} + +inline void CheckIfHeadDimsMatch(const at::Tensor& a, const at::Tensor& b) { + TORCH_CHECK(a.dim() == b.dim(), "Expected tensors to match in rank."); + auto tensor_rank = a.dim(); + TORCH_CHECK( + a.size(tensor_rank - 1) == b.size(tensor_rank - 1), + "Tensors don't match in head dim; ", + a.size(tensor_rank - 1), + " != ", + b.size(tensor_rank - 1)); +} + +inline void CheckIfBatchHeadsHeadDimMatch( + const at::Tensor& a, + const at::Tensor& b) { + TORCH_CHECK(a.dim() == b.dim(), "Expected tensors to match in rank."); + auto tensor_rank = a.dim(); + CheckIfBatchHeadsMatch(a, b); + CheckIfHeadDimsMatch(a, b); +} + +template +void CheckLogSumExp(const at::Tensor& output, const at::Tensor& logsumexp) { + // Output: [batch, *, heads, dim] + // Logsumexp: [batch, *, heads] + static_assert(NaDim >= 1 && NaDim < 4); + TORCH_CHECK( + logsumexp.scalar_type() == torch::kFloat, + "`logsumexp` must be stored in float32 data type, got ", + logsumexp.scalar_type()); + TORCH_CHECK( + logsumexp.device().is_cuda() == output.device().is_cuda(), + "Expected logsumexp to be on the same device as the operands."); + CHECK_CONTIGUOUS(logsumexp); + TORCH_CHECK( + output.dim() == NaDim + 3, + NaDim, + "-D NA expects operands to be ", + NaDim + 3, + " rank tensors, got ", + output.dim()); + TORCH_CHECK( + logsumexp.dim() == output.dim() - 1, + NaDim, + "-D NA expects logsumexp to be a ", + NaDim + 2, + " rank tensor, got ", + logsumexp.dim()); + for (size_t i = 0; i < NaDim + 2; ++i) { + TORCH_CHECK( + logsumexp.size(i) == output.size(i), + "Invalid logsumexp shape at dim ", + i, + "; " + "expected ", + output.size(i), + ", got ", + logsumexp.size(i), + "."); + } +} + +inline void CheckLogSumExpHeadsFirst( + const at::Tensor& output, + const at::Tensor& logsumexp) { + // Output: [batch, seqlen, heads, dim] + // Logsumexp: [batch, heads, seqlen] + TORCH_CHECK( + logsumexp.scalar_type() == torch::kFloat, + "`logsumexp` must be stored in float32 data type, got ", + logsumexp.scalar_type()); + TORCH_CHECK( + logsumexp.device().is_cuda() == output.device().is_cuda(), + "Expected logsumexp to be on the same device as the operands."); + CHECK_CONTIGUOUS(logsumexp); + TORCH_CHECK( + output.dim() == 4, "Expected 4D tensor, got ", output.dim(), "D."); + TORCH_CHECK( + logsumexp.dim() == 3, + "Expected 3D logsumexp tensor, got ", + logsumexp.dim(), + "D/"); + + TORCH_CHECK( + logsumexp.size(0) == output.size(0), + "Logsumexp and input tensor don't match in the batch dimension. ", + "Input tensor has batch=", + output.size(0), + ", logsumexp has batch=", + logsumexp.size(0), + "."); + + TORCH_CHECK( + logsumexp.size(1) == output.size(2), + "Logsumexp and input tensor don't match in the head dimension. ", + "Input tensor has heads=", + output.size(2), + ", logsumexp has heads=", + logsumexp.size(1), + "."); + + // NOTE: seqlen check is >= instead of ==, since LSE can be padded + TORCH_CHECK( + logsumexp.size(2) >= output.size(1), + "Logsumexp and input tensor don't match in the seqlen dimension. ", + "Input tensor has seqlen=", + output.size(1), + ", logsumexp has seqlen=", + logsumexp.size(2), + "."); +} + +inline void AssertDimsAre128BitAligned( + const at::Tensor& query, + const at::Tensor& value) { + int head_dim = static_cast(query.size(-1)); + int head_dim_value = static_cast(value.size(-1)); + + TORCH_CHECK( + query.scalar_type() == value.scalar_type(), + "Q, K, and V must match in data type, got query.dtype=", + query.scalar_type(), + ", but value.dtype=", + value.scalar_type(), + "."); + + TORCH_CHECK( + query.scalar_type() == torch::kFloat || + query.scalar_type() == torch::kFloat16 || + query.scalar_type() == torch::kBFloat16 || + query.scalar_type() == c10::ScalarType::Float8_e4m3fn || + query.scalar_type() == c10::ScalarType::Float8_e5m2, + "This NATTEN operation only supports FP32, FP16, BF16, and FP8 data types, got ", + query.scalar_type(), + "."); + + int alignment; + if (query.scalar_type() == torch::kFloat) { + alignment = 4; + } else if ( + query.scalar_type() == torch::kFloat16 || + query.scalar_type() == torch::kBFloat16) { + alignment = 8; + } else if ( + query.scalar_type() == c10::ScalarType::Float8_e4m3fn || + query.scalar_type() == c10::ScalarType::Float8_e5m2) { + alignment = 16; + } else { + TORCH_CHECK(false, "This should not have happened. Please open an issue."); + } + + TORCH_CHECK( + head_dim % alignment == 0, + "Query head dimension must be a multiple of ", + alignment, + " for this data type ", + "to meet the 128-bit alignment constraint, got ", + head_dim, + "."); + TORCH_CHECK( + head_dim_value % alignment == 0, + "Value head dimension must be a multiple of ", + alignment, + " for this data type ", + "to meet the 128-bit alignment constraint, got ", + head_dim_value, + "."); +} + +} // namespace natten diff --git a/natten/csrc/include/natten/hopper_fmha.h b/natten/csrc/include/natten/hopper_fmha.h new file mode 100644 index 00000000..00fc32a6 --- /dev/null +++ b/natten/csrc/include/natten/hopper_fmha.h @@ -0,0 +1,73 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ +/*! \file + \brief Hopper FMHA interface +*/ + +#pragma once +#include + +#include + +namespace natten { + +void hopper_fmha_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::optional& logsumexp, + bool is_causal, + float attn_scale, + int query_tile_size, + int key_tile_size, + int kernel_type, + // varlen + const at::optional& cumulative_seqlen_Q, + const at::optional& cumulative_seqlen_KV, + // only used if cumulative_seqlen_Q and cumulative_seqlen_KV are specified + int max_seqlen_Q, + int max_seqlen_KV); + +void hopper_fmha_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + bool is_causal, + float attn_scale, + int query_tile_size, + int key_tile_size, + // varlen + const at::optional& cumulative_seqlen_Q, + const at::optional& cumulative_seqlen_KV, + // only used if cumulative_seqlen_Q and cumulative_seqlen_KV are specified + int max_seqlen_Q, + int max_seqlen_KV); + +} // namespace natten diff --git a/natten/csrc/include/natten/hopper_fna.h b/natten/csrc/include/natten/hopper_fna.h new file mode 100644 index 00000000..9b31ea07 --- /dev/null +++ b/natten/csrc/include/natten/hopper_fna.h @@ -0,0 +1,155 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ +/*! \file + \brief Hopper FNA interface +*/ + +#pragma once +#include + +#include + +namespace natten { + +// Forward + +void hopper_na1d_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::optional& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& q_shape_, + const std::tuple& kv_shape_, + const std::tuple& qkv_shape_, + const std::tuple& query_tile_size, + const std::tuple& key_tile_size, + int kernel_type); + +void hopper_na2d_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::optional& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& q_shape_, + const std::tuple& kv_shape_, + const std::tuple& qkv_shape_, + const std::tuple& query_tile_size, + const std::tuple& key_tile_size, + int kernel_type); + +void hopper_na3d_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::optional& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& q_shape_, + const std::tuple& kv_shape_, + const std::tuple& qkv_shape_, + const std::tuple& query_tile_size, + const std::tuple& key_tile_size, + int kernel_type); + +// Backward + +void hopper_na1d_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& q_shape_, + const std::tuple& kv_shape_, + const std::tuple& qkv_shape_, + const std::tuple& query_tile_shape, + const std::tuple& key_tile_shape); + +void hopper_na2d_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& q_shape_, + const std::tuple& kv_shape_, + const std::tuple& qkv_shape_, + const std::tuple& query_tile_shape, + const std::tuple& key_tile_shape); + +void hopper_na3d_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& q_shape_, + const std::tuple& kv_shape_, + const std::tuple& qkv_shape_, + const std::tuple& query_tile_shape, + const std::tuple& key_tile_shape); + +} // namespace natten diff --git a/natten/csrc/include/natten/natten.h b/natten/csrc/include/natten/natten.h new file mode 100644 index 00000000..01e4ce4e --- /dev/null +++ b/natten/csrc/include/natten/natten.h @@ -0,0 +1,147 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ + +#pragma once + +#include +#include + +#if __has_include() +#include +#endif + +// kernel-builder port: upstream NATTEN gates the Hopper/Blackwell backends +// with CMake options. Here every TU derives the same setting from the CUDA +// toolkit version instead, so that the static dispatchers and the autogen +// kernel sections (which build.toml gates with `cuda-minver`) always agree. +#if !defined(NATTEN_DISABLE_HOPPER_FNA) && !defined(NATTEN_WITH_HOPPER_FNA) +#define NATTEN_WITH_HOPPER_FNA 1 +#endif +#if !defined(NATTEN_DISABLE_BLACKWELL_FNA) && \ + !defined(NATTEN_WITH_BLACKWELL_FNA) && defined(CUDART_VERSION) && \ + CUDART_VERSION >= 12080 +#define NATTEN_WITH_BLACKWELL_FNA 1 +#endif + +namespace natten { + +#define NATTEN_FAILURE(msg) throw std::runtime_error(#msg); + +#define NATTEN_CHECK(cond, msg) \ + if (!(cond)) { \ + throw std::runtime_error(#msg); \ + } + +// Causal mask helpers +template +bool any_true(TupleType v); + +template <> +inline bool any_true(std::tuple v) { + return std::get<0>(v); +} + +template <> +inline bool any_true(std::tuple v) { + return std::get<0>(v) || std::get<1>(v); +} + +template <> +inline bool any_true(std::tuple v) { + return std::get<0>(v) || std::get<1>(v) || std::get<2>(v); +} + +// Kernel size helpers +template +int32_t flatten(TupleType v); + +template <> +inline int32_t flatten(std::tuple v) { + return std::get<0>(v); +} + +template <> +inline int32_t flatten(std::tuple v) { + return std::get<0>(v) * std::get<1>(v); +} + +template <> +inline int32_t flatten(std::tuple v) { + return std::get<0>(v) * std::get<1>(v) * std::get<2>(v); +} + +template +bool all_dims_match(TupleType v); + +template <> +inline bool all_dims_match(std::tuple v) { + return true; +} + +template <> +inline bool all_dims_match(std::tuple v) { + return std::get<0>(v) == std::get<1>(v); +} + +template <> +inline bool all_dims_match(std::tuple v) { + return std::get<0>(v) == std::get<1>(v) && std::get<1>(v) == std::get<2>(v); +} + +// Runtime accessors + +template +int32_t get_from_tuple(TupleType v, size_t index); + +template <> +inline int32_t get_from_tuple(std::tuple v, size_t index) { + NATTEN_CHECK(index == 0, "Got invalid index for tuple of size 1."); + return std::get<0>(v); +} + +template <> +inline int32_t get_from_tuple(std::tuple v, size_t index) { + NATTEN_CHECK( + index == 0 || index == 1, "Got invalid index for tuple of size 2."); + if (index == 0) { + return std::get<0>(v); + } + return std::get<1>(v); +} + +template <> +inline int32_t get_from_tuple( + std::tuple v, + size_t index) { + NATTEN_CHECK( + index == 0 || index == 1 || index == 2, + "Got invalid index for tuple of size 3."); + if (index == 0) { + return std::get<0>(v); + } else if (index == 1) { + return std::get<1>(v); + } + return std::get<2>(v); +} + +} // namespace natten diff --git a/natten/csrc/include/natten/reference.h b/natten/csrc/include/natten/reference.h new file mode 100644 index 00000000..fd92b1c5 --- /dev/null +++ b/natten/csrc/include/natten/reference.h @@ -0,0 +1,134 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ +/*! \file + \brief Reference kernel interface +*/ + +#pragma once +#include + +#include + +namespace natten { + +// Forward + +void reference_na1d_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& qkv_shape, + int num_extra_kv); + +void reference_na2d_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& qkv_shape, + int num_extra_kv); + +void reference_na3d_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& qkv_shape, + int num_extra_kv); + +// Backward + +void reference_na1d_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& qkv_shape, + int num_extra_kv); + +void reference_na2d_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& qkv_shape, + int num_extra_kv); + +void reference_na3d_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& qkv_shape, + int num_extra_kv); + +} // namespace natten diff --git a/natten/csrc/include/natten/token_permute.h b/natten/csrc/include/natten/token_permute.h new file mode 100644 index 00000000..cffa54c3 --- /dev/null +++ b/natten/csrc/include/natten/token_permute.h @@ -0,0 +1,76 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ +/*! \file + \brief Token Permute / Unpermute +*/ + +#pragma once +#include + +#include + +namespace natten { + +void token_permute_1d( + at::Tensor& out, + const at::Tensor& in, + const std::tuple& tile_shape, + const std::tuple& dilation, + bool flip_tiled_dims); + +void token_permute_2d( + at::Tensor& out, + const at::Tensor& in, + const std::tuple& tile_shape, + const std::tuple& dilation, + bool flip_tiled_dims); + +void token_permute_3d( + at::Tensor& out, + const at::Tensor& in, + const std::tuple& tile_shape, + const std::tuple& dilation, + bool flip_tiled_dims); + +void token_unpermute_1d( + at::Tensor& out, + const at::Tensor& in, + const std::tuple& tile_shape, + const std::tuple& dilation, + bool flip_tiled_dims); + +void token_unpermute_2d( + at::Tensor& out, + const at::Tensor& in, + const std::tuple& tile_shape, + const std::tuple& dilation, + bool flip_tiled_dims); + +void token_unpermute_3d( + at::Tensor& out, + const at::Tensor& in, + const std::tuple& tile_shape, + const std::tuple& dilation, + bool flip_tiled_dims); + +} // namespace natten diff --git a/natten/csrc/src/blackwell_fmha.cu b/natten/csrc/src/blackwell_fmha.cu new file mode 100644 index 00000000..9104462a --- /dev/null +++ b/natten/csrc/src/blackwell_fmha.cu @@ -0,0 +1,495 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ +/*! \file + \brief Blackwell FMHA interface +*/ + +#include +#include +#include +#include +#include + +#include +#include + +#if defined(NATTEN_WITH_CUTLASS) && defined(NATTEN_WITH_BLACKWELL_FNA) +#include +#include +#include +#include +#endif + +namespace natten { + +void blackwell_fmha_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::optional& logsumexp, + bool is_causal, + float attn_scale, + int query_tile_size, + int key_tile_size, + bool run_persistent, + // varlen + const at::optional& cumulative_seqlen_Q, + const at::optional& cumulative_seqlen_KV, + // only used if cumulative_seqlen_Q and cumulative_seqlen_KV are specified + int max_seqlen_Q, + int max_seqlen_KV) { +#if defined(NATTEN_WITH_CUTLASS) && defined(NATTEN_WITH_BLACKWELL_FNA) + AssertDimsAre128BitAligned(query, value); + + CHECK_CONTIGUOUS(query); + CHECK_CONTIGUOUS(key); + CHECK_CONTIGUOUS(value); + CHECK_CONTIGUOUS(out); + + CHECK_CUDA(query); + CHECK_CUDA(key); + CHECK_CUDA(value); + CHECK_CUDA(out); + + at::cuda::OptionalCUDAGuard device_guard(query.device()); + + CheckIfPropertiesMatch(query, key, value); + CheckIfTensorShapesMatch<1>(key, value); // head_dim == head_dim_value + CheckIfTensorShapesMatch<1>(query, out); // head_dim == head_dim_value + + TORCH_CHECK(query.dim() == 4, "Tensors must be 4-D."); + TORCH_CHECK(key.dim() == 4, "Tensors must be 4-D."); + TORCH_CHECK(value.dim() == 4, "Tensors must be 4-D."); + + TORCH_CHECK( + query.size(0) == key.size(0), + "Blackwell FMHA forward: Query and key must match in batch size, got ", + "query.shape[0]=", + query.size(0), + ", key.shape[0]=", + key.size(0)); + + TORCH_CHECK( + query.size(3) == key.size(3), + "Blackwell FMHA forward: Query and key must match in head dim, got ", + "query.shape[3]=", + query.size(3), + ", key.shape[3]=", + key.size(3)); + + // GQA/MQA is supported + TORCH_CHECK( + query.size(2) >= key.size(2), + "Blackwell FMHA forward: Query heads must be greater than or equal to key/value heads, got ", + "query.shape[2]=", + query.size(2), + ", key.shape[2]=", + key.size(2)); + + TORCH_CHECK( + query.size(2) % key.size(2) == 0, + "Blackwell FMHA forward: Query heads must evenly divide key/value heads, got ", + "query.shape[2]=", + query.size(2), + ", key.shape[2]=", + key.size(2)); + + int batch_size = query.size(0); + int seqlen_q = query.size(1); + int seqlen_kv = key.size(1); + int heads_q = query.size(2); + int heads_kv = key.size(2); + int dim = query.size(3); + + if (logsumexp.has_value()) { + CheckLogSumExp<1>(out, logsumexp.value()); + CHECK_CUDA(logsumexp.value()); + } + + TORCH_CHECK( + dim > 0 && dim <= 128, + "Blackwell FMHA forward only supports head dims up to 128, got ", + dim, + "."); + if (query.scalar_type() == c10::ScalarType::Float8_e4m3fn || + query.scalar_type() == c10::ScalarType::Float8_e5m2) { + TORCH_CHECK( + dim >= 16 && dim % 16 == 0, + "Blackwell FMHA forward with FP8 requires head dims that are multiples of 16 (minimum 16), got ", + dim, + "."); + } else { + TORCH_CHECK( + dim >= 8 && dim % 8 == 0, + "Blackwell FMHA forward with FP16/BF16 requires head dims that are multiples of 8 (minimum 8), got ", + dim, + "."); + } + + cudaDeviceProp* device_props = + at::cuda::getDeviceProperties(query.device().index()); + const int cc = device_props->major * 10 + device_props->minor; + + TORCH_CHECK( + cc == 100 || cc == 103, + "Blackwell FMHA forward can only run on the Blackwell (datacenter-class) architecture (SM100, SM103)."); + + TORCH_CHECK( + query.scalar_type() == key.scalar_type() && + query.scalar_type() == value.scalar_type() && + query.scalar_type() == out.scalar_type(), + "Blackwell FMHA forward: Query, key, value, and output must match in dtype."); + + TORCH_CHECK( + query.scalar_type() == torch::kFloat16 || + query.scalar_type() == torch::kBFloat16 || + query.scalar_type() == c10::ScalarType::Float8_e4m3fn || + query.scalar_type() == c10::ScalarType::Float8_e5m2, + "Blackwell FMHA forward only supports FP16, BF16, FP8_E4M3, and FP8_E5M2."); + + // varlen + bool is_varlen = + cumulative_seqlen_Q.has_value() || cumulative_seqlen_KV.has_value(); + + void* ptr_cumulative_seqlen_Q = nullptr; + void* ptr_cumulative_seqlen_KV = nullptr; + if (is_varlen) { + TORCH_CHECK( + cumulative_seqlen_Q.has_value() && cumulative_seqlen_KV.has_value(), + "Blackwell FMHA: Both cumulative_seqlen_Q and cumulative_seqlen_KV must be specified when using varlen."); + + TORCH_CHECK( + batch_size == 1, + "Blackwell FMHA: Tensor batch size must be 1 (packed sequence layout), got ", + batch_size); + + auto& cumulative_seqlen_Q_tensor = cumulative_seqlen_Q.value(); + auto& cumulative_seqlen_KV_tensor = cumulative_seqlen_KV.value(); + + TORCH_CHECK( + cumulative_seqlen_Q_tensor.dim() == 1, + "Blackwell FMHA: cumulative_seqlen_Q is expected to be a 1-D tensor."); + TORCH_CHECK( + cumulative_seqlen_KV_tensor.dim() == 1, + "Blackwell FMHA: cumulative_seqlen_KV is expected to be a 1-D tensor."); + + TORCH_CHECK( + cumulative_seqlen_Q_tensor.size(0) == + cumulative_seqlen_KV_tensor.size(0), + "Blackwell FMHA: cumulative_seqlen_Q and cumulative_seqlen_KV must be the same size."); + + TORCH_CHECK( + cumulative_seqlen_Q_tensor.size(0) > 1, + "Blackwell FMHA: cumulative_seqlen_Q and cumulative_seqlen_KV size must be greater than 1."); + + TORCH_CHECK( + cumulative_seqlen_Q_tensor.scalar_type() == torch::kInt, + "Blackwell FMHA: cumulative_seqlen_Q is expected to be an int32 tensor, got ", + cumulative_seqlen_Q_tensor.scalar_type()); + TORCH_CHECK( + cumulative_seqlen_KV_tensor.scalar_type() == torch::kInt, + "Blackwell FMHA: cumulative_seqlen_KV is expected to be an int32 tensor, got ", + cumulative_seqlen_KV_tensor.scalar_type()); + + batch_size = cumulative_seqlen_Q_tensor.size(0) - 1; + ptr_cumulative_seqlen_Q = + static_cast(cumulative_seqlen_Q_tensor.data_ptr()); + ptr_cumulative_seqlen_KV = + static_cast(cumulative_seqlen_KV_tensor.data_ptr()); + } + // + + int device_id = query.device().index(); + auto cuda_stream = at::cuda::getCurrentCUDAStream(device_id); + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + + DISPATCH_BLACKWELL_FMHA_FORWARD( + query.scalar_type(), + dim, + query_tile_size, + key_tile_size, + run_persistent, + static_cast(query.data_ptr()), + static_cast(key.data_ptr()), + static_cast(value.data_ptr()), + static_cast(out.data_ptr()), + logsumexp.has_value() ? static_cast(logsumexp.value().data_ptr()) + : nullptr, + batch_size, + seqlen_q, + seqlen_kv, + heads_q, + heads_kv, + dim, + is_causal, + attn_scale, + // varlen parameters + is_varlen, + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // init/launch params + device_id, + cuda_stream, + query.options()); + +#else + TORCH_CHECK( + false, + "Blackwell FMHA forward: libnatten was not compiled with CUTLASS_ARCH_MMA_SM100_SUPPORTED."); +#endif +#else + TORCH_CHECK( + false, + "Blackwell FMHA forward: libnatten was not compiled for Blackwell (SM100/SM103)."); +#endif +} + +void blackwell_fmha_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + bool is_causal, + float attn_scale, + int query_tile_size, + int key_tile_size, + // varlen + const at::optional& cumulative_seqlen_Q, + const at::optional& cumulative_seqlen_KV, + // only used if cumulative_seqlen_Q and cumulative_seqlen_KV are specified + int max_seqlen_Q, + int max_seqlen_KV, + bool deterministic) { +#if defined(NATTEN_WITH_CUTLASS) && defined(NATTEN_WITH_BLACKWELL_FNA) + AssertDimsAre128BitAligned(query, value); + + CHECK_CUDA(query); + CHECK_CUDA(key); + CHECK_CUDA(value); + CHECK_CUDA(out); + CHECK_CUDA(grad_query); + CHECK_CUDA(grad_key); + CHECK_CUDA(grad_value); + CHECK_CUDA(grad_out); + CHECK_CUDA(logsumexp); + + at::cuda::OptionalCUDAGuard device_guard(query.device()); + + CHECK_CONTIGUOUS(query); + CHECK_CONTIGUOUS(key); + CHECK_CONTIGUOUS(value); + CHECK_CONTIGUOUS(grad_query); + CHECK_CONTIGUOUS(grad_key); + CHECK_CONTIGUOUS(grad_value); + CHECK_CONTIGUOUS(out); + CHECK_CONTIGUOUS(grad_out); + CHECK_CONTIGUOUS(logsumexp); + + CheckIfPropertiesMatch(query, key, value); + CheckIfPropertiesMatch(grad_value, grad_out, out); + CheckIfPropertiesMatch(grad_query, grad_key, grad_value); + CheckIfPropertiesMatch(grad_query, query, value); + + CheckIfTensorShapesMatch<1>(query, out); // head_dim == head_dim_value + CheckIfTensorShapesMatch<1>(key, value); // head_dim == head_dim_value + + CheckIfHeadDimsMatch(out, value); + CheckIfTensorShapesMatch<1>(grad_query, query); + CheckIfTensorShapesMatch<1>(grad_key, key); + CheckIfTensorShapesMatch<1>(grad_value, value); + CheckIfTensorShapesMatch<1>(grad_out, out); + + CheckLogSumExp<1>(out, logsumexp); + + TORCH_CHECK( + query.size(0) == key.size(0), + "Blackwell FMHA forward: Query and key must match in batch size, got ", + "query.shape[0]=", + query.size(0), + ", key.shape[0]=", + key.size(0)); + + TORCH_CHECK( + query.size(3) == key.size(3), + "Blackwell FMHA forward: Query and key must match in head dim, got ", + "query.shape[3]=", + query.size(3), + ", key.shape[3]=", + key.size(3)); + + // GQA/MQA is supported + TORCH_CHECK( + query.size(2) >= key.size(2), + "Blackwell FMHA forward: Query heads must be greater than or equal to key/value heads, got ", + "query.shape[2]=", + query.size(2), + ", key.shape[2]=", + key.size(2)); + + TORCH_CHECK( + query.size(2) % key.size(2) == 0, + "Blackwell FMHA forward: Query heads must evenly divide key/value heads, got ", + "query.shape[2]=", + query.size(2), + ", key.shape[2]=", + key.size(2)); + + int batch_size = query.size(0); + int seqlen_q = query.size(1); + int seqlen_kv = key.size(1); + int heads_q = query.size(2); + int heads_kv = key.size(2); + int dim = query.size(3); + + // NOTE: if FP8 backward support is ever added, alignment must be multiples + // of 16. + TORCH_CHECK( + dim > 0 && dim <= 128 && dim >= 8 && dim % 8 == 0, + "Blackwell FMHA backward requires head dims that are multiples of 8 (minimum 8, maximum 128), got ", + dim, + "."); + + TORCH_CHECK( + query.scalar_type() == torch::kFloat16 || + query.scalar_type() == torch::kBFloat16, + "Blackwell FMHA backward only supports FP16 and BF16."); + + // varlen + bool is_varlen = + cumulative_seqlen_Q.has_value() || cumulative_seqlen_KV.has_value(); + + void* ptr_cumulative_seqlen_Q = nullptr; + void* ptr_cumulative_seqlen_KV = nullptr; + if (is_varlen) { + TORCH_CHECK( + cumulative_seqlen_Q.has_value() && cumulative_seqlen_KV.has_value(), + "Blackwell FMHA: Both cumulative_seqlen_Q and cumulative_seqlen_KV must be specified when using varlen."); + + TORCH_CHECK( + batch_size == 1, + "Blackwell FMHA: Tensor batch size must be 1 (packed sequence layout), got ", + batch_size); + + auto& cumulative_seqlen_Q_tensor = cumulative_seqlen_Q.value(); + auto& cumulative_seqlen_KV_tensor = cumulative_seqlen_KV.value(); + + TORCH_CHECK( + cumulative_seqlen_Q_tensor.dim() == 1, + "Blackwell FMHA: cumulative_seqlen_Q is expected to be a 1-D tensor."); + TORCH_CHECK( + cumulative_seqlen_KV_tensor.dim() == 1, + "Blackwell FMHA: cumulative_seqlen_KV is expected to be a 1-D tensor."); + + TORCH_CHECK( + cumulative_seqlen_Q_tensor.size(0) == + cumulative_seqlen_KV_tensor.size(0), + "Blackwell FMHA: cumulative_seqlen_Q and cumulative_seqlen_KV must be the same size."); + + TORCH_CHECK( + cumulative_seqlen_Q_tensor.size(0) > 1, + "Blackwell FMHA: cumulative_seqlen_Q and cumulative_seqlen_KV size must be greater than 1."); + + TORCH_CHECK( + cumulative_seqlen_Q_tensor.scalar_type() == torch::kInt, + "Blackwell FMHA: cumulative_seqlen_Q is expected to be an int32 tensor, got ", + cumulative_seqlen_Q_tensor.scalar_type()); + TORCH_CHECK( + cumulative_seqlen_KV_tensor.scalar_type() == torch::kInt, + "Blackwell FMHA: cumulative_seqlen_KV is expected to be an int32 tensor, got ", + cumulative_seqlen_KV_tensor.scalar_type()); + + batch_size = cumulative_seqlen_Q_tensor.size(0) - 1; + ptr_cumulative_seqlen_Q = + static_cast(cumulative_seqlen_Q_tensor.data_ptr()); + ptr_cumulative_seqlen_KV = + static_cast(cumulative_seqlen_KV_tensor.data_ptr()); + } + // + + int device_id = query.device().index(); + auto cuda_stream = at::cuda::getCurrentCUDAStream(device_id); + cudaDeviceProp* device_props = at::cuda::getDeviceProperties(device_id); + const int cc = device_props->major * 10 + device_props->minor; + + TORCH_CHECK( + cc == 100 || cc == 103, + "Blackwell FMHA backward can only run on the Blackwell (datacenter-class) architecture (SM100, SM103)."); + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + + DISPATCH_BLACKWELL_FMHA_BACKWARD( + query.scalar_type(), + dim, + query_tile_size, + key_tile_size, + static_cast(query.data_ptr()), + static_cast(key.data_ptr()), + static_cast(value.data_ptr()), + static_cast(out.data_ptr()), + static_cast(logsumexp.data_ptr()), + static_cast(grad_query.data_ptr()), + static_cast(grad_key.data_ptr()), + static_cast(grad_value.data_ptr()), + static_cast(grad_out.data_ptr()), + batch_size, + seqlen_q, + seqlen_kv, + heads_q, + heads_kv, + dim, + is_causal, + attn_scale, + // varlen parameters + is_varlen, + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + deterministic, + // init/launch params + device_id, + cuda_stream, + query.options()); + +#else + TORCH_CHECK( + false, + "Blackwell FMHA backward: libnatten was not compiled with CUTLASS_ARCH_MMA_SM100_SUPPORTED."); +#endif +#else + TORCH_CHECK( + false, + "Blackwell FMHA backward: libnatten was not compiled for Blackwell (SM100/SM103)."); +#endif +} + +} // namespace natten diff --git a/natten/csrc/src/blackwell_fna_backward.cu b/natten/csrc/src/blackwell_fna_backward.cu new file mode 100644 index 00000000..ee1bd333 --- /dev/null +++ b/natten/csrc/src/blackwell_fna_backward.cu @@ -0,0 +1,406 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ +/*! \file + \brief Blackwell FNA backward Torch interface +*/ + +#include +#include +#include +#include +#include + +#include +#include + +#if defined(NATTEN_WITH_CUTLASS) && defined(NATTEN_WITH_BLACKWELL_FNA) +#include +#include + +template +auto std_tuple_to_cute_tuple(StdTuple a) { + static_assert( + std::tuple_size_v > 0 && std::tuple_size_v < 4); + + if constexpr (std::tuple_size_v == 1) { + return cute::make_tuple(std::get<0>(a)); + } else if constexpr (std::tuple_size_v == 2) { + return cute::make_tuple(std::get<0>(a), std::get<1>(a)); + } else { + return cute::make_tuple(std::get<0>(a), std::get<1>(a), std::get<2>(a)); + } +} +#endif + +namespace natten { + +template +void blackwell_fna_generic_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + const StdNADim& kernel_size, + const StdNADim& stride_, + const StdNADim& dilation_, + const StdCausal& is_causal_, + float attn_scale, + const StdNADim& q_shape_, // after token permute and padding + const StdNADim& + kv_shape_, // after token permute and padding, including extra KV + const StdNADim& + qkv_shape_, // before token permute and padding, not including extra KV + const StdNADim& query_tile_shape_, + const StdNADim& key_tile_shape_) { + static_assert( + std::tuple_size_v > 0 && std::tuple_size_v < 4); + static constexpr int kNADim = std::tuple_size_v; + static_assert(std::tuple_size_v == kNADim); + +#if defined(NATTEN_WITH_CUTLASS) && defined(NATTEN_WITH_BLACKWELL_FNA) + AssertDimsAre128BitAligned(query, value); + + CHECK_CUDA(query); + CHECK_CUDA(key); + CHECK_CUDA(value); + CHECK_CUDA(out); + CHECK_CUDA(grad_query); + CHECK_CUDA(grad_key); + CHECK_CUDA(grad_value); + CHECK_CUDA(grad_out); + CHECK_CUDA(logsumexp); + + at::cuda::OptionalCUDAGuard device_guard(query.device()); + + CHECK_CONTIGUOUS(query); + CHECK_CONTIGUOUS(key); + CHECK_CONTIGUOUS(value); + CHECK_CONTIGUOUS(grad_query); + CHECK_CONTIGUOUS(grad_key); + CHECK_CONTIGUOUS(grad_value); + CHECK_CONTIGUOUS(out); + CHECK_CONTIGUOUS(grad_out); + CHECK_CONTIGUOUS(logsumexp); + + CheckArgs(kernel_size, stride_, dilation_); + CheckArgsAgainstDim(qkv_shape_, kernel_size, dilation_); + CheckIfPropertiesMatch(query, key, value); + CheckIfPropertiesMatch(grad_value, grad_out, out); + CheckIfPropertiesMatch(grad_query, grad_key, grad_value); + CheckIfPropertiesMatch(grad_query, query, value); + + // NOTE (alih): q and kv might have slightly different shapes because we're + // padding to multiples of the tile shape. We're also supporting extra KV + // tokens, so we can't have 5D/6D tensors anymore. Seqlen mode must be + // flattened. + CheckIfTensorShapesMatch<1>(query, out); // head_dim == head_dim_value + CheckIfTensorShapesMatch<1>(key, value); // head_dim == head_dim_value + + CheckIfHeadDimsMatch(out, value); + CheckIfTensorShapesMatch<1>(grad_query, query); + CheckIfTensorShapesMatch<1>(grad_key, key); + CheckIfTensorShapesMatch<1>(grad_value, value); + CheckIfTensorShapesMatch<1>(grad_out, out); + + CheckLogSumExp<1>(out, logsumexp); + + TORCH_CHECK( + query.size(0) == key.size(0), + "Blackwell FMHA forward: Query and key must match in batch size, got ", + "query.shape[0]=", + query.size(0), + ", key.shape[0]=", + key.size(0)); + + TORCH_CHECK( + query.size(3) == key.size(3), + "Blackwell FMHA forward: Query and key must match in head dim, got ", + "query.shape[3]=", + query.size(3), + ", key.shape[3]=", + key.size(3)); + + // GQA/MQA is supported + TORCH_CHECK( + query.size(2) >= key.size(2), + "Blackwell FMHA forward: Query heads must be greater than or equal to key/value heads, got ", + "query.shape[2]=", + query.size(2), + ", key.shape[2]=", + key.size(2)); + + TORCH_CHECK( + query.size(2) % key.size(2) == 0, + "Blackwell FMHA forward: Query heads must evenly divide key/value heads, got ", + "query.shape[2]=", + query.size(2), + ", key.shape[2]=", + key.size(2)); + + int batch_size = query.size(0); + int seqlen_q = query.size(1); + int seqlen_kv = key.size(1); + int heads_q = query.size(2); + int heads_kv = key.size(2); + int dim = query.size(3); + + auto qkv_shape = std_tuple_to_cute_tuple(qkv_shape_); + auto q_shape = std_tuple_to_cute_tuple(q_shape_); + auto kv_shape = std_tuple_to_cute_tuple(kv_shape_); + + auto query_tile_shape = std_tuple_to_cute_tuple(query_tile_shape_); + auto key_tile_shape = std_tuple_to_cute_tuple(key_tile_shape_); + + auto window_size = std_tuple_to_cute_tuple(kernel_size); + auto stride = std_tuple_to_cute_tuple(stride_); + auto dilation = std_tuple_to_cute_tuple(dilation_); + auto is_causal = std_tuple_to_cute_tuple(is_causal_); + + // NOTE: if FP8 backward support is ever added, alignment must be multiples + // of 16. + TORCH_CHECK( + dim > 0 && dim <= 128 && dim >= 8 && dim % 8 == 0, + "Blackwell FNA backward requires head dims that are multiples of 8 (minimum 8, maximum 128), got ", + dim, + "."); + + TORCH_CHECK( + size(q_shape) == seqlen_q, + "Blackwell FNA backward: Q sequence length (q.shape[1]) must match the size of Q shape."); + TORCH_CHECK( + size(kv_shape) == seqlen_kv, + "Blackwell FNA backward: KV sequence length ({k,v}.shape[1]) must match the size of KV shape."); + + TORCH_CHECK( + cute::evenly_divides(q_shape, query_tile_shape) && + cute::evenly_divides(kv_shape, key_tile_shape), + "Blackwell FNA backward: Tile shapes must evenly divide input. Please pad your inputs."); + + TORCH_CHECK( + query.scalar_type() == torch::kFloat16 || + query.scalar_type() == torch::kBFloat16, + "Blackwell FNA backward only supports FP16 and BF16."); + + TORCH_CHECK( + not at::globalContext().deterministicAlgorithms(), + "Blackwell FNA backward is non-deterministic, " + "but PyTorch's deterministic mode is enabled. " + "NATTEN Python API should have avoided this; which means " + "you're probably calling the C function directly."); + + int device_id = query.device().index(); + auto cuda_stream = at::cuda::getCurrentCUDAStream(device_id); + cudaDeviceProp* device_props = at::cuda::getDeviceProperties(device_id); + const int cc = device_props->major * 10 + device_props->minor; + + TORCH_CHECK( + cc == 100 || cc == 103, + "Blackwell FMHA backward can only run on the Blackwell (datacenter-class) architecture (SM100, SM103)."); + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + + DISPATCH_BLACKWELL_FNA_BACKWARD( + kNADim, + query.scalar_type(), + dim, + is_causal, + query_tile_shape, + key_tile_shape, + static_cast(query.data_ptr()), + static_cast(key.data_ptr()), + static_cast(value.data_ptr()), + static_cast(out.data_ptr()), + static_cast(logsumexp.data_ptr()), + static_cast(grad_query.data_ptr()), + static_cast(grad_key.data_ptr()), + static_cast(grad_value.data_ptr()), + static_cast(grad_out.data_ptr()), + batch_size, + seqlen_q, + seqlen_kv, + heads_q, + heads_kv, + dim, + attn_scale, + // fna parameters + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + // init/launch params + device_id, + cuda_stream, + query.options()); + +#else + TORCH_CHECK( + false, + "Blackwell FNA backward: libnatten was not compiled with CUTLASS_ARCH_MMA_SM100_SUPPORTED."); +#endif +#else + TORCH_CHECK( + false, + "Blackwell FNA backward: libnatten was not compiled for Blackwell (SM100/SM103)."); +#endif +} + +void blackwell_na1d_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& q_shape, + const std::tuple& kv_shape, + const std::tuple& qkv_shape, + const std::tuple& query_tile_shape, + const std::tuple& key_tile_shape) { + TORCH_CHECK(query.dim() == 4, "Tensors must be 4-D."); + + blackwell_fna_generic_backward( + grad_query, + grad_key, + grad_value, + query, + key, + value, + out, + grad_out, + logsumexp, + kernel_size, + stride, + dilation, + is_causal, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + query_tile_shape, + key_tile_shape); +} + +void blackwell_na2d_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& q_shape, + const std::tuple& kv_shape, + const std::tuple& qkv_shape, + const std::tuple& query_tile_shape, + const std::tuple& key_tile_shape) { + TORCH_CHECK(query.dim() == 4, "Tensors must be 4-D."); + + blackwell_fna_generic_backward( + grad_query, + grad_key, + grad_value, + query, + key, + value, + out, + grad_out, + logsumexp, + kernel_size, + stride, + dilation, + is_causal, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + query_tile_shape, + key_tile_shape); +} + +void blackwell_na3d_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& q_shape, + const std::tuple& kv_shape, + const std::tuple& qkv_shape, + const std::tuple& query_tile_shape, + const std::tuple& key_tile_shape) { + TORCH_CHECK(query.dim() == 4, "Tensors must be 4-D."); + + blackwell_fna_generic_backward( + grad_query, + grad_key, + grad_value, + query, + key, + value, + out, + grad_out, + logsumexp, + kernel_size, + stride, + dilation, + is_causal, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + query_tile_shape, + key_tile_shape); +} + +} // namespace natten diff --git a/natten/csrc/src/blackwell_fna_forward.cu b/natten/csrc/src/blackwell_fna_forward.cu new file mode 100644 index 00000000..b0ede04d --- /dev/null +++ b/natten/csrc/src/blackwell_fna_forward.cu @@ -0,0 +1,383 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ +/*! \file + \brief Blackwell FNA Torch interface +*/ + +#include +#include +#include +#include +#include + +#include +#include + +#if defined(NATTEN_WITH_CUTLASS) && defined(NATTEN_WITH_BLACKWELL_FNA) +#include +#include + +template +auto std_tuple_to_cute_tuple(StdTuple a) { + static_assert( + std::tuple_size_v > 0 && std::tuple_size_v < 4); + + if constexpr (std::tuple_size_v == 1) { + return cute::make_tuple(std::get<0>(a)); + } else if constexpr (std::tuple_size_v == 2) { + return cute::make_tuple(std::get<0>(a), std::get<1>(a)); + } else { + return cute::make_tuple(std::get<0>(a), std::get<1>(a), std::get<2>(a)); + } +} +#endif + +namespace natten { + +template +void blackwell_fna_generic_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::optional& logsumexp, + const StdNADim& kernel_size, + const StdNADim& stride_, + const StdNADim& dilation_, + const StdCausal& is_causal_, + float attn_scale, + const StdNADim& q_shape_, // after token permute and padding + const StdNADim& kv_shape_, // after token permute and padding + const StdNADim& qkv_shape_, // before token permute and padding + const StdNADim& query_tile_shape_, + const StdNADim& key_tile_shape_, + bool run_persistent) { + static_assert( + std::tuple_size_v > 0 && std::tuple_size_v < 4); + static constexpr int kNADim = std::tuple_size_v; + static_assert(std::tuple_size_v == kNADim); + +#if defined(NATTEN_WITH_CUTLASS) && defined(NATTEN_WITH_BLACKWELL_FNA) + AssertDimsAre128BitAligned(query, value); + + CHECK_CONTIGUOUS(query); + CHECK_CONTIGUOUS(key); + CHECK_CONTIGUOUS(value); + CHECK_CONTIGUOUS(out); + + CHECK_CUDA(query); + CHECK_CUDA(key); + CHECK_CUDA(value); + CHECK_CUDA(out); + + at::cuda::OptionalCUDAGuard device_guard(query.device()); + + CheckArgs(kernel_size, stride_, dilation_); + CheckIfPropertiesMatch(query, key, value); + + // NOTE (alih): q and kv might have slightly different shapes because we're + // padding to multiples of the tile shape. Seqlen mode must be flattened. + CheckIfTensorShapesMatch<1>(key, value); // head_dim == head_dim_value + CheckIfTensorShapesMatch<1>(query, out); // head_dim == head_dim_value + + TORCH_CHECK(query.dim() == 4, "Tensors must be 4-D."); + TORCH_CHECK(key.dim() == 4, "Tensors must be 4-D."); + TORCH_CHECK(value.dim() == 4, "Tensors must be 4-D."); + + TORCH_CHECK( + query.size(0) == key.size(0), + "Blackwell FMHA forward: Query and key must match in batch size, got ", + "query.shape[0]=", + query.size(0), + ", key.shape[0]=", + key.size(0)); + + TORCH_CHECK( + query.size(3) == key.size(3), + "Blackwell FMHA forward: Query and key must match in head dim, got ", + "query.shape[3]=", + query.size(3), + ", key.shape[3]=", + key.size(3)); + + // GQA/MQA is supported + TORCH_CHECK( + query.size(2) >= key.size(2), + "Blackwell FMHA forward: Query heads must be greater than or equal to key/value heads, got ", + "query.shape[2]=", + query.size(2), + ", key.shape[2]=", + key.size(2)); + + TORCH_CHECK( + query.size(2) % key.size(2) == 0, + "Blackwell FMHA forward: Query heads must evenly divide key/value heads, got ", + "query.shape[2]=", + query.size(2), + ", key.shape[2]=", + key.size(2)); + + int batch_size = query.size(0); + int seqlen_q = query.size(1); + int seqlen_kv = key.size(1); + int heads_q = query.size(2); + int heads_kv = key.size(2); + int dim = query.size(3); + + if (logsumexp.has_value()) { + CheckLogSumExp<1>(out, logsumexp.value()); + CHECK_CUDA(logsumexp.value()); + } + + CheckArgsAgainstDim(qkv_shape_, kernel_size, dilation_); + + auto qkv_shape = std_tuple_to_cute_tuple(qkv_shape_); + auto q_shape = std_tuple_to_cute_tuple(q_shape_); + auto kv_shape = std_tuple_to_cute_tuple(kv_shape_); + + auto query_tile_shape = std_tuple_to_cute_tuple(query_tile_shape_); + auto key_tile_shape = std_tuple_to_cute_tuple(key_tile_shape_); + + auto window_size = std_tuple_to_cute_tuple(kernel_size); + auto stride = std_tuple_to_cute_tuple(stride_); + auto dilation = std_tuple_to_cute_tuple(dilation_); + auto is_causal = std_tuple_to_cute_tuple(is_causal_); + + TORCH_CHECK( + size(q_shape) == seqlen_q, + "Blackwell FNA forward: Q sequence length (q.shape[1]) must match the size of QKV shape."); + TORCH_CHECK( + size(kv_shape) == seqlen_kv, + "Blackwell FNA forward: KV sequence length ({k,v}.shape[1]) must match the size of QKV shape."); + + TORCH_CHECK( + dim > 0 && dim <= 128, + "Blackwell FNA forward only supports head dims up to 128, got ", + dim, + "."); + if (query.scalar_type() == c10::ScalarType::Float8_e4m3fn || + query.scalar_type() == c10::ScalarType::Float8_e5m2) { + TORCH_CHECK( + dim >= 16 && dim % 16 == 0, + "Blackwell FNA forward with FP8 requires head dims that are multiples of 16 (minimum 16), got ", + dim, + "."); + } else { + TORCH_CHECK( + dim >= 8 && dim % 8 == 0, + "Blackwell FNA forward with FP16/BF16 requires head dims that are multiples of 8 (minimum 8), got ", + dim, + "."); + } + + cudaDeviceProp* device_props = + at::cuda::getDeviceProperties(query.device().index()); + const int cc = device_props->major * 10 + device_props->minor; + + TORCH_CHECK( + cc == 100 || cc == 103, + "Blackwell FNA forward can only run on the Blackwell (datacenter-class) architecture (SM100, SM103)."); + + TORCH_CHECK( + query.scalar_type() == key.scalar_type() && + query.scalar_type() == value.scalar_type() && + query.scalar_type() == out.scalar_type(), + "Blackwell FNA forward: Query, key, value, and output must match in dtype."); + + TORCH_CHECK( + query.scalar_type() == torch::kFloat16 || + query.scalar_type() == torch::kBFloat16 || + query.scalar_type() == c10::ScalarType::Float8_e4m3fn || + query.scalar_type() == c10::ScalarType::Float8_e5m2, + "Blackwell FNA forward only supports FP16, BF16, FP8_E4M3, and FP8_E5M2."); + + int device_id = query.device().index(); + auto cuda_stream = at::cuda::getCurrentCUDAStream(device_id); + + TORCH_CHECK( + cute::evenly_divides(q_shape, query_tile_shape) && + cute::evenly_divides(kv_shape, key_tile_shape), + "Blackwell FNA forward: Tile shapes must evenly divide input. Please pad your inputs."); + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + + DISPATCH_BLACKWELL_FNA_FORWARD( + kNADim, + query.scalar_type(), + dim, + is_causal, + query_tile_shape, + key_tile_shape, + run_persistent, + static_cast(query.data_ptr()), + static_cast(key.data_ptr()), + static_cast(value.data_ptr()), + static_cast(out.data_ptr()), + logsumexp.has_value() ? static_cast(logsumexp.value().data_ptr()) + : nullptr, + batch_size, + seqlen_q, + seqlen_kv, + heads_q, + heads_kv, + dim, + attn_scale, + // fna / fusion parameters + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + // init/launch params + device_id, + cuda_stream, + query.options()); + +#else + TORCH_CHECK( + false, + "Blackwell FNA forward: libnatten was not compiled with CUTLASS_ARCH_MMA_SM100_SUPPORTED."); +#endif +#else + TORCH_CHECK( + false, + "Blackwell FNA forward: libnatten was not compiled for Blackwell (SM100/SM103)."); +#endif +} + +void blackwell_na1d_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::optional& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& q_shape, + const std::tuple& kv_shape, + const std::tuple& qkv_shape, + const std::tuple& query_tile_shape, + const std::tuple& key_tile_shape, + bool run_persistent) { + TORCH_CHECK(query.dim() == 4, "Tensors must be 4-D."); + + blackwell_fna_generic_forward( + out, + query, + key, + value, + logsumexp, + kernel_size, + stride, + dilation, + is_causal, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + query_tile_shape, + key_tile_shape, + run_persistent); +} + +void blackwell_na2d_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::optional& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& q_shape, + const std::tuple& kv_shape, + const std::tuple& qkv_shape, + const std::tuple& query_tile_shape, + const std::tuple& key_tile_shape, + bool run_persistent) { + TORCH_CHECK(query.dim() == 4, "Tensors must be 4-D."); + + blackwell_fna_generic_forward( + out, + query, + key, + value, + logsumexp, + kernel_size, + stride, + dilation, + is_causal, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + query_tile_shape, + key_tile_shape, + run_persistent); +} + +void blackwell_na3d_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::optional& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& q_shape, + const std::tuple& kv_shape, + const std::tuple& qkv_shape, + const std::tuple& query_tile_shape, + const std::tuple& key_tile_shape, + bool run_persistent) { + TORCH_CHECK(query.dim() == 4, "Tensors must be 4-D."); + + blackwell_fna_generic_forward( + out, + query, + key, + value, + logsumexp, + kernel_size, + stride, + dilation, + is_causal, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + query_tile_shape, + key_tile_shape, + run_persistent); +} + +} // namespace natten diff --git a/natten/csrc/src/compute_delta.cu b/natten/csrc/src/compute_delta.cu new file mode 100644 index 00000000..02c548d1 --- /dev/null +++ b/natten/csrc/src/compute_delta.cu @@ -0,0 +1,123 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ +/*! \file + \brief compute_delta interface + mostly used to test the kernel. +*/ + +#include +#include +#include +#include + +#include +#include +#include + +#ifdef NATTEN_WITH_CUTLASS +#include +#include +#include +#include +#endif + +namespace natten { + +void compute_delta( + const at::Tensor& out, + const at::Tensor& d_out, + at::Tensor& delta) { +#ifdef NATTEN_WITH_CUTLASS + at::cuda::OptionalCUDAGuard device_guard(out.device()); + + CHECK_CONTIGUOUS(out); + CHECK_CONTIGUOUS(d_out); + CHECK_CONTIGUOUS(delta); + CHECK_CUDA(out); + CHECK_CUDA(d_out); + CHECK_CUDA(delta); + + TORCH_CHECK(out.dim() == 4, "out must be 4-D."); + TORCH_CHECK(d_out.dim() == 4, "d_out must be 4-D."); + TORCH_CHECK(delta.dim() == 3, "delta must be 3-D."); + + CheckIfTensorShapesMatch<1>(out, d_out); + CheckLogSumExp<1>(out, delta); + + int batch = out.size(0); + int seqlen_Q = out.size(1); + int heads = out.size(2); + int dim = out.size(3); + + NATTEN_CHECK( + out.scalar_type() == torch::kFloat || + out.scalar_type() == torch::kFloat16 || + out.scalar_type() == torch::kBFloat16, + "`compute_delta` only supports FP32, FP16 and BF16 operands."); + NATTEN_CHECK( + d_out.scalar_type() == out.scalar_type(), + "`compute_delta` input operands must match in dtype."); + NATTEN_CHECK( + delta.scalar_type() == torch::kFloat, "Delta is always in fp32."); + + if (out.scalar_type() == torch::kFloat) { + natten::cuda::compute_delta( + at::cuda::getCurrentCUDAStream(out.device().index()), + static_cast(out.data_ptr()), + static_cast(d_out.data_ptr()), + static_cast(delta.data_ptr()), + batch, + heads, + seqlen_Q, + dim); + } else if (out.scalar_type() == torch::kFloat16) { + natten::cuda::compute_delta( + at::cuda::getCurrentCUDAStream(out.device().index()), + static_cast(out.data_ptr()), + static_cast(d_out.data_ptr()), + static_cast(delta.data_ptr()), + batch, + heads, + seqlen_Q, + dim); + } else if (out.scalar_type() == torch::kBFloat16) { + natten::cuda::compute_delta( + at::cuda::getCurrentCUDAStream(out.device().index()), + static_cast(out.data_ptr()), + static_cast(d_out.data_ptr()), + static_cast(delta.data_ptr()), + batch, + heads, + seqlen_Q, + dim); + } else { + NATTEN_FAILURE("`compute_delta` is unavailable for this element type."); + } + +#else + NATTEN_FAILURE( + "`compute_delta` is only available when NATTEN is built with CUTLASS."); +#endif +} + +} // namespace natten diff --git a/natten/csrc/src/fmha.cu b/natten/csrc/src/fmha.cu new file mode 100644 index 00000000..1cdc05bf --- /dev/null +++ b/natten/csrc/src/fmha.cu @@ -0,0 +1,424 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ +/*! \file + \brief FMHA interface +*/ + +#include +#include +#include +#include +#include + +#include +#include +#include + +#ifdef NATTEN_WITH_CUTLASS +#include +#include +#include +#endif + +namespace natten { + +void fmha_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::optional& logsumexp, + bool is_causal, + float attn_scale, + int query_tile_size, + int key_tile_size, + // varlen + const at::optional& cumulative_seqlen_Q, + const at::optional& cumulative_seqlen_KV, + // only used if cumulative_seqlen_Q and cumulative_seqlen_KV are specified + int max_seqlen_Q, + int max_seqlen_KV) { +#ifdef NATTEN_WITH_CUTLASS + AssertDimsAre128BitAligned(query, value); + + CHECK_CONTIGUOUS(query); + CHECK_CONTIGUOUS(key); + CHECK_CONTIGUOUS(value); + CHECK_CONTIGUOUS(out); + + CHECK_CUDA(query); + CHECK_CUDA(key); + CHECK_CUDA(value); + CHECK_CUDA(out); + + at::cuda::OptionalCUDAGuard device_guard(query.device()); + + CheckIfPropertiesMatch(query, key, value); + CheckIfTensorShapesMatchExceptHeadDim<1>(key, value); + CheckIfTensorShapesMatchExceptHeadDim<1>(query, out); + CheckIfBatchHeadsHeadDimMatch(query, key); + CheckIfHeadDimsMatch(out, value); + + TORCH_CHECK(query.dim() == 4, "Tensors must be 4-D."); + TORCH_CHECK(key.dim() == 4, "Tensors must be 4-D."); + TORCH_CHECK(value.dim() == 4, "Tensors must be 4-D."); + + int batch_size = query.size(0); + int seqlen_q = query.size(1); + int seqlen_kv = key.size(1); + int heads = query.size(2); + int dim = query.size(3); + int dim_value = value.size(3); + + if (logsumexp.has_value()) { + CheckLogSumExp<1>(out, logsumexp.value()); + CHECK_CUDA(logsumexp.value()); + } + + // varlen + bool is_varlen = + cumulative_seqlen_Q.has_value() || cumulative_seqlen_KV.has_value(); + + void* ptr_cumulative_seqlen_Q = nullptr; + void* ptr_cumulative_seqlen_KV = nullptr; + if (is_varlen) { + TORCH_CHECK( + cumulative_seqlen_Q.has_value() && cumulative_seqlen_KV.has_value(), + "CUTLASS FMHA: Both cumulative_seqlen_Q and cumulative_seqlen_KV must be specified when using varlen."); + + TORCH_CHECK( + batch_size == 1, + "CUTLASS FMHA: Tensor batch size must be 1 (packed sequence layout), got ", + batch_size); + + auto& cumulative_seqlen_Q_tensor = cumulative_seqlen_Q.value(); + auto& cumulative_seqlen_KV_tensor = cumulative_seqlen_KV.value(); + + TORCH_CHECK( + cumulative_seqlen_Q_tensor.dim() == 1, + "CUTLASS FMHA: cumulative_seqlen_Q is expected to be a 1-D tensor."); + TORCH_CHECK( + cumulative_seqlen_KV_tensor.dim() == 1, + "CUTLASS FMHA: cumulative_seqlen_KV is expected to be a 1-D tensor."); + + TORCH_CHECK( + cumulative_seqlen_Q_tensor.size(0) == + cumulative_seqlen_KV_tensor.size(0), + "CUTLASS FMHA: cumulative_seqlen_Q and cumulative_seqlen_KV must be the same size."); + + TORCH_CHECK( + cumulative_seqlen_Q_tensor.size(0) > 1, + "CUTLASS FMHA: cumulative_seqlen_Q and cumulative_seqlen_KV size must be greater than 1."); + + TORCH_CHECK( + cumulative_seqlen_Q_tensor.scalar_type() == torch::kInt, + "CUTLASS FMHA: cumulative_seqlen_Q is expected to be an int32 tensor, got ", + cumulative_seqlen_Q_tensor.scalar_type()); + TORCH_CHECK( + cumulative_seqlen_KV_tensor.scalar_type() == torch::kInt, + "CUTLASS FMHA: cumulative_seqlen_KV is expected to be an int32 tensor, got ", + cumulative_seqlen_KV_tensor.scalar_type()); + + batch_size = cumulative_seqlen_Q_tensor.size(0) - 1; + ptr_cumulative_seqlen_Q = + static_cast(cumulative_seqlen_Q_tensor.data_ptr()); + ptr_cumulative_seqlen_KV = + static_cast(cumulative_seqlen_KV_tensor.data_ptr()); + } + // + + at::Tensor workspace; + auto alloc_bytes = [&workspace, &query]( + void** ptr, int64_t bytes, bool zfill) { + workspace = at::empty({bytes}, query.options().dtype(at::ScalarType::Byte)); + if (zfill) { + workspace.zero_(); + } + *ptr = static_cast(workspace.data_ptr()); + }; + + cudaDeviceProp* device_props = + at::cuda::getDeviceProperties(query.device().index()); + const int cc = device_props->major * 10 + device_props->minor; + const size_t max_smem = device_props->sharedMemPerBlockOptin; + + if (cc >= 80 || (cc >= 50 && query.scalar_type() != torch::kBFloat16)) { + natten::cuda::fmha::fmha_forward_generic( + query.scalar_type(), + cc, + max_smem, + at::cuda::getCurrentCUDAStream(query.device().index()), + alloc_bytes, + static_cast(query.data_ptr()), + static_cast(key.data_ptr()), + static_cast(value.data_ptr()), + static_cast(out.data_ptr()), + batch_size, + seqlen_q, + seqlen_kv, + heads, + dim, + dim_value, + is_causal, + attn_scale, + logsumexp.has_value() ? static_cast(logsumexp.value().data_ptr()) + : nullptr, + // varlen parameters + is_varlen, + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // init/launch params + query_tile_size, + key_tile_size); + } else { + NATTEN_FAILURE( + "FMHA kernels are only available on devices with " + "compute capability >= 50 for FP16/FP32 inputs, and devices with " + "compute capability >= 80 for FP32, BF16, and FP16 inputs."); + } +#else + TORCH_CHECK(false, "libnatten not compiled with CUTLASS."); +#endif +} + +void fmha_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + bool is_causal, + float attn_scale, + int query_tile_size, + int key_tile_size, + int num_splits_key, + bool compute_delta_with_torch, + // varlen + const at::optional& cumulative_seqlen_Q, + const at::optional& cumulative_seqlen_KV, + // only used if cumulative_seqlen_Q and cumulative_seqlen_KV are specified + int max_seqlen_Q, + int max_seqlen_KV) { +#ifdef NATTEN_WITH_CUTLASS + AssertDimsAre128BitAligned(query, value); + + CHECK_CUDA(query); + CHECK_CUDA(key); + CHECK_CUDA(value); + CHECK_CUDA(out); + CHECK_CUDA(grad_query); + CHECK_CUDA(grad_key); + CHECK_CUDA(grad_value); + CHECK_CUDA(grad_out); + CHECK_CUDA(logsumexp); + + at::cuda::OptionalCUDAGuard device_guard(query.device()); + + CHECK_CONTIGUOUS(query); + CHECK_CONTIGUOUS(key); + CHECK_CONTIGUOUS(value); + CHECK_CONTIGUOUS(grad_query); + CHECK_CONTIGUOUS(grad_key); + CHECK_CONTIGUOUS(grad_value); + CHECK_CONTIGUOUS(out); + CHECK_CONTIGUOUS(grad_out); + CHECK_CONTIGUOUS(logsumexp); + + CheckIfPropertiesMatch(query, key, value); + CheckIfPropertiesMatch(grad_value, grad_out, out); + CheckIfPropertiesMatch(grad_query, grad_key, grad_value); + CheckIfPropertiesMatch(grad_query, query, value); + + CheckIfTensorShapesMatchExceptHeadDim<1>(query, out); + CheckIfTensorShapesMatchExceptHeadDim<1>(key, value); + CheckIfBatchHeadsHeadDimMatch(query, key); + CheckIfHeadDimsMatch(out, value); + CheckIfTensorShapesMatch<1>(grad_query, query); + CheckIfTensorShapesMatch<1>(grad_key, key); + CheckIfTensorShapesMatch<1>(grad_value, value); + CheckIfTensorShapesMatch<1>(grad_out, out); + + CheckLogSumExp<1>(out, logsumexp); + + int batch_size = query.size(0); + int seqlen_q = query.size(1); + int seqlen_kv = key.size(1); + int heads = query.size(2); + int dim = query.size(3); + int dim_value = value.size(3); + + at::Tensor workspace; + auto alloc_bytes = [&workspace, &query]( + void** ptr, int64_t bytes, bool zfill) { + workspace = at::empty({bytes}, query.options().dtype(at::ScalarType::Byte)); + if (zfill) { + workspace.zero_(); + } + *ptr = static_cast(workspace.data_ptr()); + }; + at::Tensor delta; + if (compute_delta_with_torch) { + delta = (grad_out.to(at::kFloat) * out.to(at::kFloat)).sum(-1); + } else { + delta = torch::empty( + {batch_size, seqlen_q, heads}, query.options().dtype(at::kFloat)); + compute_delta(out, grad_out, delta); + } + TORCH_CHECK(delta.size(0) == batch_size); + TORCH_CHECK(delta.size(1) == seqlen_q); + TORCH_CHECK(delta.size(2) == heads); + + if (at::globalContext().deterministicAlgorithms()) { + TORCH_CHECK( + not compute_delta_with_torch, + "Computing delta with PyTorch is not guaranteed to be deterministic!"); + TORCH_CHECK( + num_splits_key <= 1, + "CUTLASS FMHA backward was called with KV parallelism, " + "which makes it algorithm non-deterministic, " + "but PyTorch's deterministic mode is enabled. " + "NATTEN Python API should have avoided this; which means " + "you're probably calling the C function directly."); + } + + // varlen + bool is_varlen = + cumulative_seqlen_Q.has_value() || cumulative_seqlen_KV.has_value(); + + void* ptr_cumulative_seqlen_Q = nullptr; + void* ptr_cumulative_seqlen_KV = nullptr; + if (is_varlen) { + TORCH_CHECK( + cumulative_seqlen_Q.has_value() && cumulative_seqlen_KV.has_value(), + "CUTLASS FMHA: Both cumulative_seqlen_Q and cumulative_seqlen_KV must be specified when using varlen."); + + TORCH_CHECK( + batch_size == 1, + "CUTLASS FMHA: Tensor batch size must be 1 (packed sequence layout), got ", + batch_size); + + auto& cumulative_seqlen_Q_tensor = cumulative_seqlen_Q.value(); + auto& cumulative_seqlen_KV_tensor = cumulative_seqlen_KV.value(); + + TORCH_CHECK( + cumulative_seqlen_Q_tensor.dim() == 1, + "CUTLASS FMHA: cumulative_seqlen_Q is expected to be a 1-D tensor."); + TORCH_CHECK( + cumulative_seqlen_KV_tensor.dim() == 1, + "CUTLASS FMHA: cumulative_seqlen_KV is expected to be a 1-D tensor."); + + TORCH_CHECK( + cumulative_seqlen_Q_tensor.size(0) == + cumulative_seqlen_KV_tensor.size(0), + "CUTLASS FMHA: cumulative_seqlen_Q and cumulative_seqlen_KV must be the same size."); + + TORCH_CHECK( + cumulative_seqlen_Q_tensor.size(0) > 1, + "CUTLASS FMHA: cumulative_seqlen_Q and cumulative_seqlen_KV size must be greater than 1."); + + TORCH_CHECK( + cumulative_seqlen_Q_tensor.scalar_type() == torch::kInt, + "CUTLASS FMHA: cumulative_seqlen_Q is expected to be an int32 tensor, got ", + cumulative_seqlen_Q_tensor.scalar_type()); + TORCH_CHECK( + cumulative_seqlen_KV_tensor.scalar_type() == torch::kInt, + "CUTLASS FMHA: cumulative_seqlen_KV is expected to be an int32 tensor, got ", + cumulative_seqlen_KV_tensor.scalar_type()); + + batch_size = cumulative_seqlen_Q_tensor.size(0) - 1; + ptr_cumulative_seqlen_Q = + static_cast(cumulative_seqlen_Q_tensor.data_ptr()); + ptr_cumulative_seqlen_KV = + static_cast(cumulative_seqlen_KV_tensor.data_ptr()); + } + // + + auto seqlen_kv_ = is_varlen ? max_seqlen_KV : seqlen_kv; + auto num_kv_tiles = (seqlen_kv_ + key_tile_size - 1) / key_tile_size; + TORCH_CHECK( + num_splits_key <= num_kv_tiles, + "CUTLASS FMHA backward: backward_kv_splits must be less than or equal to ", + "KV sequence length divided by tile size, got backward_kv_splits=", + num_splits_key, + ", ceil_div(seqlen_kv, key_tile_size)=", + num_kv_tiles, + "."); + + cudaDeviceProp* device_props = + at::cuda::getDeviceProperties(query.device().index()); + const int cc = device_props->major * 10 + device_props->minor; + const size_t max_smem = device_props->sharedMemPerBlockOptin; + + if (cc >= 80 || (cc >= 50 && query.scalar_type() != torch::kBFloat16)) { + natten::cuda::fmha::fmha_backward_generic( + query.scalar_type(), + cc, + max_smem, + at::cuda::getCurrentCUDAStream(query.device().index()), + alloc_bytes, + static_cast(grad_out.data_ptr()), + static_cast(query.data_ptr()), + static_cast(key.data_ptr()), + static_cast(value.data_ptr()), + static_cast(logsumexp.data_ptr()), + static_cast(delta.data_ptr()), + static_cast(out.data_ptr()), + static_cast(grad_query.data_ptr()), + static_cast(grad_key.data_ptr()), + static_cast(grad_value.data_ptr()), + batch_size, + seqlen_q, + seqlen_kv, + heads, + dim, + dim_value, + is_causal, + attn_scale, + // varlen parameters + is_varlen, + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // init/launch params + query_tile_size, + key_tile_size, + num_splits_key); + } else { + NATTEN_FAILURE( + "FMHA are only available on devices with " + "compute capability >= 50 for FP16/FP32 inputs, and devices with " + "compute capability >= 80 for FP32, BF16, and FP16 inputs."); + } +#else + TORCH_CHECK(false, "libnatten not compiled with CUTLASS."); +#endif +} + +} // namespace natten diff --git a/natten/csrc/src/fna_backward.cu b/natten/csrc/src/fna_backward.cu new file mode 100644 index 00000000..edd60f94 --- /dev/null +++ b/natten/csrc/src/fna_backward.cu @@ -0,0 +1,345 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ +/*! \file + \brief FNA backward interface +*/ + +#include +#include +#include +#include +#include + +#include +#include +#include + +#ifdef NATTEN_WITH_CUTLASS +#include +#include +#endif + +namespace natten { + +template +auto tuple_product(StdTuple a) { + static_assert( + std::tuple_size_v > 0 && std::tuple_size_v < 4); + + if constexpr (std::tuple_size_v == 1) { + return std::get<0>(a); + } else if constexpr (std::tuple_size_v == 2) { + return std::get<0>(a) * std::get<1>(a); + } else { + return std::get<0>(a) * std::get<1>(a) * std::get<2>(a); + } +} + +template +void fna_generic_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + const StdNADim& kernel_size, + const StdNADim& stride, + const StdNADim& dilation, + const StdCausal& is_causal, + float attn_scale, + const StdNADim& qkv_shape, + const StdNADim& query_tile_size, + const StdNADim& key_tile_size, + const StdNADim& num_splits_key, + bool compute_delta_with_torch) { + static_assert( + std::tuple_size_v > 0 && std::tuple_size_v < 4); + static constexpr int kNADim = std::tuple_size_v; + static_assert(std::tuple_size_v == kNADim); + +#ifdef NATTEN_WITH_CUTLASS + AssertDimsAre128BitAligned(query, value); + + CHECK_CUDA(query); + CHECK_CUDA(key); + CHECK_CUDA(value); + CHECK_CUDA(out); + CHECK_CUDA(grad_query); + CHECK_CUDA(grad_key); + CHECK_CUDA(grad_value); + CHECK_CUDA(grad_out); + CHECK_CUDA(logsumexp); + + at::cuda::OptionalCUDAGuard device_guard(query.device()); + + CHECK_CONTIGUOUS(query); + CHECK_CONTIGUOUS(key); + CHECK_CONTIGUOUS(value); + CHECK_CONTIGUOUS(grad_query); + CHECK_CONTIGUOUS(grad_key); + CHECK_CONTIGUOUS(grad_value); + CHECK_CONTIGUOUS(out); + CHECK_CONTIGUOUS(grad_out); + CHECK_CONTIGUOUS(logsumexp); + + CheckArgs(kernel_size, stride, dilation); + CheckIfPropertiesMatch(query, key, value); + CheckIfPropertiesMatch(grad_value, grad_out, out); + CheckIfPropertiesMatch(grad_query, grad_key, grad_value); + CheckIfPropertiesMatch(grad_query, query, value); + + CheckIfTensorShapesMatch(query, key); + CheckIfTensorShapesMatchExceptHeadDim(query, value); + CheckIfTensorShapesMatch(out, value); + CheckIfTensorShapesMatch(grad_query, query); + CheckIfTensorShapesMatch(grad_key, key); + CheckIfTensorShapesMatch(grad_value, value); + CheckIfTensorShapesMatch(grad_out, out); + + CheckLogSumExp(out, logsumexp); + + int batch_size = query.size(0); + int heads = query.size(kNADim + 1); + int dim = query.size(kNADim + 2); + int dim_value = value.size(kNADim + 2); + auto seqlen = tuple_product(qkv_shape); + CheckArgsAgainstDim(qkv_shape, kernel_size, dilation); + + at::Tensor workspace; + auto alloc_bytes = [&workspace, &query]( + void** ptr, int64_t bytes, bool zfill) { + workspace = at::empty({bytes}, query.options().dtype(at::ScalarType::Byte)); + if (zfill) { + workspace.zero_(); + } + *ptr = static_cast(workspace.data_ptr()); + }; + at::Tensor delta; + if (compute_delta_with_torch) { + delta = (grad_out.to(at::kFloat) * out.to(at::kFloat)) + .flatten(1, kNADim) + .sum(-1); + } else { + delta = torch::empty( + {batch_size, seqlen, heads}, query.options().dtype(at::kFloat)); + auto out_ = torch::flatten(out, /*start_dim=*/1, /*end_dim=*/kNADim); + auto grad_out_ = + torch::flatten(grad_out, /*start_dim=*/1, /*end_dim=*/kNADim); + compute_delta(out_, grad_out_, delta); + } + TORCH_CHECK(delta.size(0) == batch_size); + TORCH_CHECK(delta.size(1) == seqlen); + TORCH_CHECK(delta.size(2) == heads); + if (at::globalContext().deterministicAlgorithms()) { + TORCH_CHECK( + not compute_delta_with_torch, + "Computing delta with PyTorch is not guaranteed to be deterministic!"); + TORCH_CHECK( + natten::flatten(num_splits_key) <= 1, + "FNA-backward was called with KV parallelism, " + "which makes it algorithm non-deterministic, " + "but PyTorch's deterministic mode is enabled. " + "NATTEN Python API should have avoided this; which means " + "you're probably calling the C function directly."); + } + + cudaDeviceProp* device_props = + at::cuda::getDeviceProperties(query.device().index()); + const int cc = device_props->major * 10 + device_props->minor; + const size_t max_smem = device_props->sharedMemPerBlockOptin; + + if (cc >= 80 || (cc >= 50 && query.scalar_type() != torch::kBFloat16)) { + natten::cuda::fna::fna_backward_generic( + query.scalar_type(), + cc, + max_smem, + at::cuda::getCurrentCUDAStream(query.device().index()), + alloc_bytes, + static_cast(grad_out.data_ptr()), + static_cast(query.data_ptr()), + static_cast(key.data_ptr()), + static_cast(value.data_ptr()), + static_cast(logsumexp.data_ptr()), + static_cast(delta.data_ptr()), + static_cast(out.data_ptr()), + static_cast(grad_query.data_ptr()), + static_cast(grad_key.data_ptr()), + static_cast(grad_value.data_ptr()), + batch_size, + qkv_shape, + heads, + dim, + dim_value, + kernel_size, + stride, + dilation, + is_causal, + attn_scale, + query_tile_size, + key_tile_size, + num_splits_key); + } else { + NATTEN_FAILURE( + "Fused kernels are only available on devices with " + "compute capability >= 50 for FP16/FP32 inputs, and devices with " + "compute capability >= 80 for FP32, BF16, and FP16 inputs."); + } +#else + TORCH_CHECK(false, "libnatten not compiled with CUTLASS."); +#endif +} + +void na1d_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& query_tile_size, + const std::tuple& key_tile_size, + const std::tuple& num_splits_key, + bool compute_delta_with_torch) { + TORCH_CHECK(query.dim() == 4, "Tensors must be 4-D."); + + fna_generic_backward( + grad_query, + grad_key, + grad_value, + query, + key, + value, + out, + grad_out, + logsumexp, + kernel_size, + stride, + dilation, + is_causal, + attn_scale, + {query.size(1)}, + query_tile_size, + key_tile_size, + num_splits_key, + compute_delta_with_torch); +} + +void na2d_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& query_tile_size, + const std::tuple& key_tile_size, + const std::tuple& num_splits_key, + bool compute_delta_with_torch) { + TORCH_CHECK(query.dim() == 5, "Tensors must be 5-D."); + + fna_generic_backward( + grad_query, + grad_key, + grad_value, + query, + key, + value, + out, + grad_out, + logsumexp, + kernel_size, + stride, + dilation, + is_causal, + attn_scale, + {query.size(1), query.size(2)}, + query_tile_size, + key_tile_size, + num_splits_key, + compute_delta_with_torch); +} + +void na3d_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& query_tile_size, + const std::tuple& key_tile_size, + const std::tuple& num_splits_key, + bool compute_delta_with_torch) { + TORCH_CHECK(query.dim() == 6, "Tensors must be 6-D."); + + fna_generic_backward( + grad_query, + grad_key, + grad_value, + query, + key, + value, + out, + grad_out, + logsumexp, + kernel_size, + stride, + dilation, + is_causal, + attn_scale, + {query.size(1), query.size(2), query.size(3)}, + query_tile_size, + key_tile_size, + num_splits_key, + compute_delta_with_torch); +} + +} // namespace natten diff --git a/natten/csrc/src/fna_forward.cu b/natten/csrc/src/fna_forward.cu new file mode 100644 index 00000000..265ffe83 --- /dev/null +++ b/natten/csrc/src/fna_forward.cu @@ -0,0 +1,238 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ +/*! \file + \brief FNA forward interface +*/ + +#include +#include +#include +#include +#include + +#include +#include + +#ifdef NATTEN_WITH_CUTLASS +#include +#include +#endif + +namespace natten { + +template +void fna_generic_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::optional& logsumexp, + const StdNADim& kernel_size, + const StdNADim& stride, + const StdNADim& dilation, + const StdCausal& is_causal, + float attn_scale, + const StdNADim& qkv_shape, + const StdNADim& query_tile_size, + const StdNADim& key_tile_size) { + static_assert( + std::tuple_size_v > 0 && std::tuple_size_v < 4); + static constexpr int kNADim = std::tuple_size_v; + static_assert(std::tuple_size_v == kNADim); + +#ifdef NATTEN_WITH_CUTLASS + AssertDimsAre128BitAligned(query, value); + + CHECK_CONTIGUOUS(query); + CHECK_CONTIGUOUS(key); + CHECK_CONTIGUOUS(value); + CHECK_CONTIGUOUS(out); + + CHECK_CUDA(query); + CHECK_CUDA(key); + CHECK_CUDA(value); + CHECK_CUDA(out); + + at::cuda::OptionalCUDAGuard device_guard(query.device()); + + CheckArgs(kernel_size, stride, dilation); + CheckIfPropertiesMatch(query, key, value); + CheckIfTensorShapesMatch(query, key); + CheckIfTensorShapesMatchExceptHeadDim(query, value); + CheckIfTensorShapesMatch(out, value); + + int batch_size = query.size(0); + int heads = query.size(kNADim + 1); + int dim = query.size(kNADim + 2); + int dim_value = value.size(kNADim + 2); + CheckArgsAgainstDim(qkv_shape, kernel_size, dilation); + if (logsumexp.has_value()) { + CheckLogSumExp(out, logsumexp.value()); + CHECK_CUDA(logsumexp.value()); + } + + at::Tensor workspace; + auto alloc_bytes = [&workspace, &query]( + void** ptr, int64_t bytes, bool zfill) { + workspace = at::empty({bytes}, query.options().dtype(at::ScalarType::Byte)); + if (zfill) { + workspace.zero_(); + } + *ptr = static_cast(workspace.data_ptr()); + }; + + cudaDeviceProp* device_props = + at::cuda::getDeviceProperties(query.device().index()); + const int cc = device_props->major * 10 + device_props->minor; + const size_t max_smem = device_props->sharedMemPerBlockOptin; + + if (cc >= 80 || (cc >= 50 && query.scalar_type() != torch::kBFloat16)) { + natten::cuda::fna::fna_forward_generic( + query.scalar_type(), + cc, + max_smem, + at::cuda::getCurrentCUDAStream(query.device().index()), + alloc_bytes, + static_cast(query.data_ptr()), + static_cast(key.data_ptr()), + static_cast(value.data_ptr()), + static_cast(out.data_ptr()), + batch_size, + qkv_shape, + heads, + dim, + dim_value, + kernel_size, + stride, + dilation, + is_causal, + attn_scale, + logsumexp.has_value() ? static_cast(logsumexp.value().data_ptr()) + : nullptr, + query_tile_size, + key_tile_size); + } else { + NATTEN_FAILURE( + "Fused kernels are only available on devices with " + "compute capability >= 50 for FP16/FP32 inputs, and devices with " + "compute capability >= 80 for FP32, BF16, and FP16 inputs."); + } +#else + TORCH_CHECK(false, "libnatten not compiled with CUTLASS."); +#endif +} + +void na1d_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::optional& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& query_tile_size, + const std::tuple& key_tile_size) { + TORCH_CHECK(query.dim() == 4, "Tensors must be 4-D."); + + fna_generic_forward( + out, + query, + key, + value, + logsumexp, + kernel_size, + stride, + dilation, + is_causal, + attn_scale, + {query.size(1)}, + query_tile_size, + key_tile_size); +} + +void na2d_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::optional& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& query_tile_size, + const std::tuple& key_tile_size) { + TORCH_CHECK(query.dim() == 5, "Tensors must be 5-D."); + + fna_generic_forward( + out, + query, + key, + value, + logsumexp, + kernel_size, + stride, + dilation, + is_causal, + attn_scale, + {query.size(1), query.size(2)}, + query_tile_size, + key_tile_size); +} + +void na3d_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::optional& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& query_tile_size, + const std::tuple& key_tile_size) { + TORCH_CHECK(query.dim() == 6, "Tensors must be 6-D."); + + fna_generic_forward( + out, + query, + key, + value, + logsumexp, + kernel_size, + stride, + dilation, + is_causal, + attn_scale, + {query.size(1), query.size(2), query.size(3)}, + query_tile_size, + key_tile_size); +} + +} // namespace natten diff --git a/natten/csrc/src/hopper_fmha.cu b/natten/csrc/src/hopper_fmha.cu new file mode 100644 index 00000000..ad674a45 --- /dev/null +++ b/natten/csrc/src/hopper_fmha.cu @@ -0,0 +1,409 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ +/*! \file + \brief Hopper FMHA interface +*/ + +#include +#include +#include +#include +#include + +#include +#include + +#include +#if defined(NATTEN_WITH_CUTLASS) && defined(NATTEN_WITH_HOPPER_FNA) +#include +#include +#endif + +namespace natten { + +void hopper_fmha_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::optional& logsumexp, + bool is_causal, + float attn_scale, + int query_tile_size, + int key_tile_size, + int kernel_type, + // varlen + const at::optional& cumulative_seqlen_Q, + const at::optional& cumulative_seqlen_KV, + // only used if cumulative_seqlen_Q and cumulative_seqlen_KV are specified + int max_seqlen_Q, + int max_seqlen_KV) { +#if defined(NATTEN_WITH_CUTLASS) && defined(NATTEN_WITH_HOPPER_FNA) + AssertDimsAre128BitAligned(query, value); + + CHECK_CONTIGUOUS(query); + CHECK_CONTIGUOUS(key); + CHECK_CONTIGUOUS(value); + CHECK_CONTIGUOUS(out); + + CHECK_CUDA(query); + CHECK_CUDA(key); + CHECK_CUDA(value); + CHECK_CUDA(out); + + at::cuda::OptionalCUDAGuard device_guard(query.device()); + + CheckIfPropertiesMatch(query, key, value); + CheckIfTensorShapesMatch<1>(key, value); + CheckIfTensorShapesMatch<1>(query, out); + CheckIfBatchHeadsHeadDimMatch(query, key); + + TORCH_CHECK(query.dim() == 4, "Tensors must be 4-D."); + TORCH_CHECK(key.dim() == 4, "Tensors must be 4-D."); + TORCH_CHECK(value.dim() == 4, "Tensors must be 4-D."); + + int batch_size = query.size(0); + int seqlen_q = query.size(1); + int seqlen_kv = key.size(1); + int heads = query.size(2); + int dim = query.size(3); + + if (logsumexp.has_value()) { + CheckLogSumExp<1>(out, logsumexp.value()); + CHECK_CUDA(logsumexp.value()); + } + + TORCH_CHECK( + dim == 32 || dim == 64 || dim == 128 || dim == 256, + "Hopper FMHA only supports head dims 32, 64, 128, and 256 for now."); + + cudaDeviceProp* device_props = + at::cuda::getDeviceProperties(query.device().index()); + const int cc = device_props->major * 10 + device_props->minor; + TORCH_CHECK( + cc == 90, + "This operation can only run on the Hopper architecture (SM90)."); + + TORCH_CHECK( + query.scalar_type() == key.scalar_type() && + query.scalar_type() == value.scalar_type() && + query.scalar_type() == out.scalar_type(), + "Query, key, value, and output must match in dtype."); + + TORCH_CHECK( + query.scalar_type() == torch::kFloat16 || + query.scalar_type() == torch::kBFloat16, + "Only FP16/BF16 is supported for now."); + + // varlen + bool is_varlen = + cumulative_seqlen_Q.has_value() || cumulative_seqlen_KV.has_value(); + + void* ptr_cumulative_seqlen_Q = nullptr; + void* ptr_cumulative_seqlen_KV = nullptr; + if (is_varlen) { + TORCH_CHECK( + cumulative_seqlen_Q.has_value() && cumulative_seqlen_KV.has_value(), + "Hopper FMHA: Both cumulative_seqlen_Q and cumulative_seqlen_KV must be specified when using varlen."); + + TORCH_CHECK( + batch_size == 1, + "Hopper FMHA: Tensor batch size must be 1 (packed sequence layout), got ", + batch_size); + + auto& cumulative_seqlen_Q_tensor = cumulative_seqlen_Q.value(); + auto& cumulative_seqlen_KV_tensor = cumulative_seqlen_KV.value(); + + TORCH_CHECK( + cumulative_seqlen_Q_tensor.dim() == 1, + "Hopper FMHA: cumulative_seqlen_Q is expected to be a 1-D tensor."); + TORCH_CHECK( + cumulative_seqlen_KV_tensor.dim() == 1, + "Hopper FMHA: cumulative_seqlen_KV is expected to be a 1-D tensor."); + + TORCH_CHECK( + cumulative_seqlen_Q_tensor.size(0) == + cumulative_seqlen_KV_tensor.size(0), + "Hopper FMHA: cumulative_seqlen_Q and cumulative_seqlen_KV must be the same size."); + + TORCH_CHECK( + cumulative_seqlen_Q_tensor.size(0) > 1, + "Hopper FMHA: cumulative_seqlen_Q and cumulative_seqlen_KV size must be greater than 1."); + + TORCH_CHECK( + cumulative_seqlen_Q_tensor.scalar_type() == torch::kInt, + "Hopper FMHA: cumulative_seqlen_Q is expected to be an int32 tensor, got ", + cumulative_seqlen_Q_tensor.scalar_type()); + TORCH_CHECK( + cumulative_seqlen_KV_tensor.scalar_type() == torch::kInt, + "Hopper FMHA: cumulative_seqlen_KV is expected to be an int32 tensor, got ", + cumulative_seqlen_KV_tensor.scalar_type()); + + batch_size = cumulative_seqlen_Q_tensor.size(0) - 1; + ptr_cumulative_seqlen_Q = + static_cast(cumulative_seqlen_Q_tensor.data_ptr()); + ptr_cumulative_seqlen_KV = + static_cast(cumulative_seqlen_KV_tensor.data_ptr()); + } + // + + int device_id = query.device().index(); + auto cuda_stream = at::cuda::getCurrentCUDAStream(device_id); + + auto kernel_type_ = + natten::cuda::hopper::kernel_type_int_to_enum_type(kernel_type); + TORCH_CHECK( + kernel_type_ != natten::cuda::hopper::HopperKernelSchedule::Invalid, + "Got invalid kernel_type argument."); + +#if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) + + DISPATCH_HOPPER_FMHA_FORWARD( + query.scalar_type(), + dim, + query_tile_size, + key_tile_size, + kernel_type_, + static_cast(query.data_ptr()), + static_cast(key.data_ptr()), + static_cast(value.data_ptr()), + static_cast(out.data_ptr()), + logsumexp.has_value() ? static_cast(logsumexp.value().data_ptr()) + : nullptr, + batch_size, + seqlen_q, + seqlen_kv, + heads, + dim, + is_causal, + attn_scale, + // varlen parameters + is_varlen, + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // init/launch params + device_id, + cuda_stream, + query.options()); + +#else + TORCH_CHECK( + false, + "libnatten was not compiled with CUTLASS_ARCH_MMA_SM90_SUPPORTED."); +#endif +#else + TORCH_CHECK(false, "libnatten was not compiled for Hopper (SM90)."); +#endif +} + +void hopper_fmha_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + bool is_causal, + float attn_scale, + int query_tile_size, + int key_tile_size, + // varlen + const at::optional& cumulative_seqlen_Q, + const at::optional& cumulative_seqlen_KV, + // only used if cumulative_seqlen_Q and cumulative_seqlen_KV are specified + int max_seqlen_Q, + int max_seqlen_KV) { +#if defined(NATTEN_WITH_CUTLASS) && defined(NATTEN_WITH_HOPPER_FNA) + AssertDimsAre128BitAligned(query, value); + + CHECK_CUDA(query); + CHECK_CUDA(key); + CHECK_CUDA(value); + CHECK_CUDA(out); + CHECK_CUDA(grad_query); + CHECK_CUDA(grad_key); + CHECK_CUDA(grad_value); + CHECK_CUDA(grad_out); + CHECK_CUDA(logsumexp); + + at::cuda::OptionalCUDAGuard device_guard(query.device()); + + CHECK_CONTIGUOUS(query); + CHECK_CONTIGUOUS(key); + CHECK_CONTIGUOUS(value); + CHECK_CONTIGUOUS(grad_query); + CHECK_CONTIGUOUS(grad_key); + CHECK_CONTIGUOUS(grad_value); + CHECK_CONTIGUOUS(out); + CHECK_CONTIGUOUS(grad_out); + CHECK_CONTIGUOUS(logsumexp); + + CheckIfPropertiesMatch(query, key, value); + CheckIfPropertiesMatch(grad_value, grad_out, out); + CheckIfPropertiesMatch(grad_query, grad_key, grad_value); + CheckIfPropertiesMatch(grad_query, query, value); + + CheckIfTensorShapesMatch<1>(query, out); + CheckIfTensorShapesMatch<1>(key, value); + CheckIfBatchHeadsHeadDimMatch(query, key); + CheckIfHeadDimsMatch(out, value); + CheckIfTensorShapesMatch<1>(grad_query, query); + CheckIfTensorShapesMatch<1>(grad_key, key); + CheckIfTensorShapesMatch<1>(grad_value, value); + CheckIfTensorShapesMatch<1>(grad_out, out); + + CheckLogSumExp<1>(out, logsumexp); + + int batch_size = query.size(0); + int seqlen_q = query.size(1); + int seqlen_kv = key.size(1); + int heads = query.size(2); + int dim = query.size(3); + + TORCH_CHECK( + dim == 32 || dim == 64 || dim == 128, + "Hopper FMHA backward pass only supports head dims 32, 64, and 128 for now."); + + TORCH_CHECK( + query.scalar_type() == torch::kFloat16 || + query.scalar_type() == torch::kBFloat16, + "Only FP16/BF16 is supported for now."); + + TORCH_CHECK( + not at::globalContext().deterministicAlgorithms(), + "Hopper FMHA backward pass is non-deterministic, " + "but PyTorch's deterministic mode is enabled. " + "NATTEN Python API should have avoided this; which means " + "you're probably calling the C function directly."); + + // varlen + bool is_varlen = + cumulative_seqlen_Q.has_value() || cumulative_seqlen_KV.has_value(); + + void* ptr_cumulative_seqlen_Q = nullptr; + void* ptr_cumulative_seqlen_KV = nullptr; + if (is_varlen) { + TORCH_CHECK( + cumulative_seqlen_Q.has_value() && cumulative_seqlen_KV.has_value(), + "Hopper FMHA backward: Both cumulative_seqlen_Q and cumulative_seqlen_KV must be specified when using varlen."); + + TORCH_CHECK( + batch_size == 1, + "Hopper FMHA backward: Tensor batch size must be 1 (packed sequence layout), got ", + batch_size); + + auto& cumulative_seqlen_Q_tensor = cumulative_seqlen_Q.value(); + auto& cumulative_seqlen_KV_tensor = cumulative_seqlen_KV.value(); + + TORCH_CHECK( + cumulative_seqlen_Q_tensor.dim() == 1, + "Hopper FMHA backward: cumulative_seqlen_Q is expected to be a 1-D tensor."); + TORCH_CHECK( + cumulative_seqlen_KV_tensor.dim() == 1, + "Hopper FMHA backward: cumulative_seqlen_KV is expected to be a 1-D tensor."); + + TORCH_CHECK( + cumulative_seqlen_Q_tensor.size(0) == + cumulative_seqlen_KV_tensor.size(0), + "Hopper FMHA backward: cumulative_seqlen_Q and cumulative_seqlen_KV must be the same size."); + + TORCH_CHECK( + cumulative_seqlen_Q_tensor.size(0) > 1, + "Hopper FMHA backward: cumulative_seqlen_Q and cumulative_seqlen_KV size must be greater than 1."); + + TORCH_CHECK( + cumulative_seqlen_Q_tensor.scalar_type() == torch::kInt, + "Hopper FMHA backward: cumulative_seqlen_Q is expected to be an int32 tensor, got ", + cumulative_seqlen_Q_tensor.scalar_type()); + TORCH_CHECK( + cumulative_seqlen_KV_tensor.scalar_type() == torch::kInt, + "Hopper FMHA backward: cumulative_seqlen_KV is expected to be an int32 tensor, got ", + cumulative_seqlen_KV_tensor.scalar_type()); + + batch_size = cumulative_seqlen_Q_tensor.size(0) - 1; + ptr_cumulative_seqlen_Q = + static_cast(cumulative_seqlen_Q_tensor.data_ptr()); + ptr_cumulative_seqlen_KV = + static_cast(cumulative_seqlen_KV_tensor.data_ptr()); + } + // + + int device_id = query.device().index(); + auto cuda_stream = at::cuda::getCurrentCUDAStream(device_id); + cudaDeviceProp* device_props = at::cuda::getDeviceProperties(device_id); + const int cc = device_props->major * 10 + device_props->minor; + + TORCH_CHECK( + cc == 90, + "This operation can only run on the Hopper architecture (SM90)."); + +#if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) + + DISPATCH_HOPPER_FMHA_BACKWARD( + query.scalar_type(), + dim, + query_tile_size, + key_tile_size, + static_cast(query.data_ptr()), + static_cast(key.data_ptr()), + static_cast(value.data_ptr()), + static_cast(out.data_ptr()), + static_cast(logsumexp.data_ptr()), + static_cast(grad_query.data_ptr()), + static_cast(grad_key.data_ptr()), + static_cast(grad_value.data_ptr()), + static_cast(grad_out.data_ptr()), + batch_size, + seqlen_q, + seqlen_kv, + heads, + dim, + is_causal, + attn_scale, + // varlen parameters + is_varlen, + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // init/launch params + device_id, + cuda_stream, + query.options()); + +#else + TORCH_CHECK( + false, + "libnatten was not compiled with CUTLASS_ARCH_MMA_SM90_SUPPORTED."); +#endif +#else + TORCH_CHECK(false, "libnatten was not compiled for Hopper (SM90)."); +#endif +} + +} // namespace natten diff --git a/natten/csrc/src/hopper_fna_backward.cu b/natten/csrc/src/hopper_fna_backward.cu new file mode 100644 index 00000000..d04e668a --- /dev/null +++ b/natten/csrc/src/hopper_fna_backward.cu @@ -0,0 +1,358 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ +/*! \file + \brief Hopper FNA backward Torch interface +*/ + +#include +#include +#include +#include +#include + +#include +#include + +#if defined(NATTEN_WITH_CUTLASS) && defined(NATTEN_WITH_HOPPER_FNA) +#include +#include + +template +auto std_tuple_to_cute_tuple(StdTuple a) { + static_assert( + std::tuple_size_v > 0 && std::tuple_size_v < 4); + + if constexpr (std::tuple_size_v == 1) { + return cute::make_tuple(std::get<0>(a)); + } else if constexpr (std::tuple_size_v == 2) { + return cute::make_tuple(std::get<0>(a), std::get<1>(a)); + } else { + return cute::make_tuple(std::get<0>(a), std::get<1>(a), std::get<2>(a)); + } +} +#endif + +namespace natten { + +template +void hopper_fna_generic_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + const StdNADim& kernel_size, + const StdNADim& stride_, + const StdNADim& dilation_, + const StdCausal& is_causal_, + float attn_scale, + const StdNADim& q_shape_, // after token permute and padding + const StdNADim& kv_shape_, // after token permute and padding + const StdNADim& qkv_shape_, // before token permute and padding + const StdNADim& query_tile_shape_, + const StdNADim& key_tile_shape_) { + static_assert( + std::tuple_size_v > 0 && std::tuple_size_v < 4); + static constexpr int kNADim = std::tuple_size_v; + static_assert(std::tuple_size_v == kNADim); + +#if defined(NATTEN_WITH_CUTLASS) && defined(NATTEN_WITH_HOPPER_FNA) + AssertDimsAre128BitAligned(query, value); + + CHECK_CONTIGUOUS(query); + CHECK_CONTIGUOUS(key); + CHECK_CONTIGUOUS(value); + CHECK_CONTIGUOUS(grad_query); + CHECK_CONTIGUOUS(grad_key); + CHECK_CONTIGUOUS(grad_value); + CHECK_CONTIGUOUS(out); + CHECK_CONTIGUOUS(grad_out); + CHECK_CONTIGUOUS(logsumexp); + + CHECK_CUDA(query); + CHECK_CUDA(key); + CHECK_CUDA(value); + CHECK_CUDA(out); + CHECK_CUDA(grad_query); + CHECK_CUDA(grad_key); + CHECK_CUDA(grad_value); + CHECK_CUDA(grad_out); + CHECK_CUDA(logsumexp); + + at::cuda::OptionalCUDAGuard device_guard(query.device()); + + CheckArgs(kernel_size, stride_, dilation_); + CheckIfPropertiesMatch(query, key, value); + CheckIfPropertiesMatch(grad_value, grad_out, out); + CheckIfPropertiesMatch(grad_query, grad_key, grad_value); + CheckIfPropertiesMatch(grad_query, query, value); + + // NOTE (alih): q and kv might have slightly different shapes because we're + // padding to multiples of the tile shape. Seqlen mode must be flattened. + CheckIfTensorShapesMatch<1>(query, out); + CheckIfTensorShapesMatch<1>(key, value); + CheckIfTensorShapesMatch<1>(query, grad_query); + CheckIfTensorShapesMatch<1>(key, grad_key); + CheckIfTensorShapesMatch<1>(value, grad_value); + CheckIfTensorShapesMatch<1>(out, grad_out); + + int batch_size = query.size(0); + int seqlen_q = query.size(1); + int seqlen_kv = key.size(1); + int heads = query.size(2); + int dim = query.size(3); + + CheckArgsAgainstDim(qkv_shape_, kernel_size, dilation_); + + CheckLogSumExp<1>(out, logsumexp); + + auto qkv_shape = std_tuple_to_cute_tuple(qkv_shape_); + auto q_shape = std_tuple_to_cute_tuple(q_shape_); + auto kv_shape = std_tuple_to_cute_tuple(kv_shape_); + + auto query_tile_shape = std_tuple_to_cute_tuple(query_tile_shape_); + auto key_tile_shape = std_tuple_to_cute_tuple(key_tile_shape_); + + auto window_size = std_tuple_to_cute_tuple(kernel_size); + auto stride = std_tuple_to_cute_tuple(stride_); + auto dilation = std_tuple_to_cute_tuple(dilation_); + auto is_causal = std_tuple_to_cute_tuple(is_causal_); + + TORCH_CHECK( + size(q_shape) == seqlen_q, + "Q's sequence length (q.shape[1]) must match the size of Q shape."); + TORCH_CHECK( + size(kv_shape) == seqlen_kv, + "KV's sequence length ({k,v}.shape[1]) must match the size of KV shape."); + + TORCH_CHECK( + cute::evenly_divides(q_shape, query_tile_shape) && + cute::evenly_divides(kv_shape, key_tile_shape), + "Tile shapes must evenly divide input. Please pad your inputs."); + + TORCH_CHECK( + query.scalar_type() == torch::kFloat16 || + query.scalar_type() == torch::kBFloat16, + "Only FP16/BF16 is supported for now."); + + TORCH_CHECK( + dim == 32 || dim == 64 || dim == 128, + "FNA Hopper only supports head dims 32, 64, and 128."); + + TORCH_CHECK( + not at::globalContext().deterministicAlgorithms(), + "Hopper FNA backward pass is non-deterministic, " + "but PyTorch's deterministic mode is enabled. " + "NATTEN Python API should have avoided this; which means " + "you're probably calling the C function directly."); + + int device_id = query.device().index(); + auto cuda_stream = at::cuda::getCurrentCUDAStream(device_id); + cudaDeviceProp* device_props = at::cuda::getDeviceProperties(device_id); + const int cc = device_props->major * 10 + device_props->minor; + + TORCH_CHECK( + cc == 90, + "This operation can only run on the Hopper architecture (SM90)."); + +#if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) + + DISPATCH_HOPPER_FNA_BACKWARD( + kNADim, + query.scalar_type(), + dim, + is_causal, + query_tile_shape, + key_tile_shape, + static_cast(query.data_ptr()), + static_cast(key.data_ptr()), + static_cast(value.data_ptr()), + static_cast(out.data_ptr()), + static_cast(logsumexp.data_ptr()), + static_cast(grad_query.data_ptr()), + static_cast(grad_key.data_ptr()), + static_cast(grad_value.data_ptr()), + static_cast(grad_out.data_ptr()), + batch_size, + seqlen_q, + seqlen_kv, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale, + cuda_stream, + query.options()); + +#else + TORCH_CHECK( + false, + "libnatten was not compiled with CUTLASS_ARCH_MMA_SM90_SUPPORTED."); +#endif +#else + TORCH_CHECK(false, "libnatten was not compiled for Hopper (SM90)."); +#endif +} + +void hopper_na1d_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& q_shape, + const std::tuple& kv_shape, + const std::tuple& qkv_shape, + const std::tuple& query_tile_shape, + const std::tuple& key_tile_shape) { + TORCH_CHECK(query.dim() == 4, "Tensors must be 4-D."); + + hopper_fna_generic_backward( + grad_query, + grad_key, + grad_value, + query, + key, + value, + out, + grad_out, + logsumexp, + kernel_size, + stride, + dilation, + is_causal, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + query_tile_shape, + key_tile_shape); +} + +void hopper_na2d_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& q_shape, + const std::tuple& kv_shape, + const std::tuple& qkv_shape, + const std::tuple& query_tile_shape, + const std::tuple& key_tile_shape) { + TORCH_CHECK(query.dim() == 4, "Tensors must be 4-D."); + + hopper_fna_generic_backward( + grad_query, + grad_key, + grad_value, + query, + key, + value, + out, + grad_out, + logsumexp, + kernel_size, + stride, + dilation, + is_causal, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + query_tile_shape, + key_tile_shape); +} + +void hopper_na3d_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& q_shape, + const std::tuple& kv_shape, + const std::tuple& qkv_shape, + const std::tuple& query_tile_shape, + const std::tuple& key_tile_shape) { + TORCH_CHECK(query.dim() == 4, "Tensors must be 4-D."); + + hopper_fna_generic_backward( + grad_query, + grad_key, + grad_value, + query, + key, + value, + out, + grad_out, + logsumexp, + kernel_size, + stride, + dilation, + is_causal, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + query_tile_shape, + key_tile_shape); +} + +} // namespace natten diff --git a/natten/csrc/src/hopper_fna_forward.cu b/natten/csrc/src/hopper_fna_forward.cu new file mode 100644 index 00000000..19e5bd18 --- /dev/null +++ b/natten/csrc/src/hopper_fna_forward.cu @@ -0,0 +1,327 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ +/*! \file + \brief Hopper FNA Torch interface +*/ + +#include +#include +#include +#include +#include + +#include +#include + +#include +#if defined(NATTEN_WITH_CUTLASS) && defined(NATTEN_WITH_HOPPER_FNA) +#include +#include + +template +auto std_tuple_to_cute_tuple(StdTuple a) { + static_assert( + std::tuple_size_v > 0 && std::tuple_size_v < 4); + + if constexpr (std::tuple_size_v == 1) { + return cute::make_tuple(std::get<0>(a)); + } else if constexpr (std::tuple_size_v == 2) { + return cute::make_tuple(std::get<0>(a), std::get<1>(a)); + } else { + return cute::make_tuple(std::get<0>(a), std::get<1>(a), std::get<2>(a)); + } +} +#endif + +namespace natten { + +template +void hopper_fna_generic_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::optional& logsumexp, + const StdNADim& kernel_size, + const StdNADim& stride_, + const StdNADim& dilation_, + const StdCausal& is_causal_, + float attn_scale, + const StdNADim& q_shape_, // after token permute and padding + const StdNADim& kv_shape_, // after token permute and padding + const StdNADim& qkv_shape_, // before token permute and padding + const StdNADim& query_tile_shape_, + const StdNADim& key_tile_shape_, + int kernel_type) { + static_assert( + std::tuple_size_v > 0 && std::tuple_size_v < 4); + static constexpr int kNADim = std::tuple_size_v; + static_assert(std::tuple_size_v == kNADim); + +#if defined(NATTEN_WITH_CUTLASS) && defined(NATTEN_WITH_HOPPER_FNA) + AssertDimsAre128BitAligned(query, value); + + CHECK_CONTIGUOUS(query); + CHECK_CONTIGUOUS(key); + CHECK_CONTIGUOUS(value); + CHECK_CONTIGUOUS(out); + + CHECK_CUDA(query); + CHECK_CUDA(key); + CHECK_CUDA(value); + CHECK_CUDA(out); + + at::cuda::OptionalCUDAGuard device_guard(query.device()); + + CheckArgs(kernel_size, stride_, dilation_); + CheckIfPropertiesMatch(query, key, value); + + // NOTE (alih): q and kv might have slightly different shapes because we're + // padding to multiples of the tile shape. Seqlen mode must be flattened. + CheckIfTensorShapesMatch<1>(query, out); + CheckIfTensorShapesMatch<1>(key, value); + + TORCH_CHECK(query.dim() == 4, "Tensors must be 4-D."); + TORCH_CHECK(key.dim() == 4, "Tensors must be 4-D."); + TORCH_CHECK(value.dim() == 4, "Tensors must be 4-D."); + + int batch_size = query.size(0); + int seqlen_q = query.size(1); + int seqlen_kv = key.size(1); + int heads = query.size(2); + int dim = query.size(3); + + CheckArgsAgainstDim(qkv_shape_, kernel_size, dilation_); + + auto qkv_shape = std_tuple_to_cute_tuple(qkv_shape_); + auto q_shape = std_tuple_to_cute_tuple(q_shape_); + auto kv_shape = std_tuple_to_cute_tuple(kv_shape_); + + auto query_tile_shape = std_tuple_to_cute_tuple(query_tile_shape_); + auto key_tile_shape = std_tuple_to_cute_tuple(key_tile_shape_); + + auto window_size = std_tuple_to_cute_tuple(kernel_size); + auto stride = std_tuple_to_cute_tuple(stride_); + auto dilation = std_tuple_to_cute_tuple(dilation_); + auto is_causal = std_tuple_to_cute_tuple(is_causal_); + + TORCH_CHECK( + size(q_shape) == seqlen_q, + "Q's sequence length (q.shape[1]) must match the size of Q shape."); + TORCH_CHECK( + size(kv_shape) == seqlen_kv, + "KV's sequence length ({k,v}.shape[1]) must match the size of KV shape."); + + if (logsumexp.has_value()) { + CheckLogSumExp<1>(out, logsumexp.value()); + CHECK_CUDA(logsumexp.value()); + } + + TORCH_CHECK( + dim == 32 || dim == 64 || dim == 128 || dim == 256, + "Hopper FNA only supports head dims 32, 64, 128, and 256 for now."); + + cudaDeviceProp* device_props = + at::cuda::getDeviceProperties(query.device().index()); + const int cc = device_props->major * 10 + device_props->minor; + TORCH_CHECK( + cc == 90, + "This operation can only run on the Hopper architecture (SM90)."); + + TORCH_CHECK( + query.scalar_type() == key.scalar_type() && + query.scalar_type() == value.scalar_type() && + query.scalar_type() == out.scalar_type(), + "Query, key, value, and output must match in dtype."); + + TORCH_CHECK( + query.scalar_type() == torch::kFloat16 || + query.scalar_type() == torch::kBFloat16, + "Only FP16/BF16 is supported for now."); + + int device_id = query.device().index(); + auto cuda_stream = at::cuda::getCurrentCUDAStream(device_id); + + auto kernel_type_ = + natten::cuda::hopper::kernel_type_int_to_enum_type(kernel_type); + TORCH_CHECK( + kernel_type_ != natten::cuda::hopper::HopperKernelSchedule::Invalid, + "Got invalid kernel_type argument."); + +#if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) + + DISPATCH_HOPPER_FNA_FORWARD( + kNADim, + query.scalar_type(), + dim, + is_causal, + query_tile_shape, + key_tile_shape, + kernel_type_, + static_cast(query.data_ptr()), + static_cast(key.data_ptr()), + static_cast(value.data_ptr()), + static_cast(out.data_ptr()), + logsumexp.has_value() ? static_cast(logsumexp.value().data_ptr()) + : nullptr, + batch_size, + seqlen_q, + seqlen_kv, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale, + cuda_stream, + query.options()); + +#else + TORCH_CHECK( + false, + "libnatten was not compiled with CUTLASS_ARCH_MMA_SM90_SUPPORTED."); +#endif +#else + TORCH_CHECK(false, "libnatten was not compiled for Hopper (SM90)."); +#endif +} + +void hopper_na1d_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::optional& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& q_shape, + const std::tuple& kv_shape, + const std::tuple& qkv_shape, + const std::tuple& query_tile_shape, + const std::tuple& key_tile_shape, + int kernel_type) { + TORCH_CHECK(query.dim() == 4, "Tensors must be 4-D."); + + hopper_fna_generic_forward( + out, + query, + key, + value, + logsumexp, + kernel_size, + stride, + dilation, + is_causal, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + query_tile_shape, + key_tile_shape, + kernel_type); +} + +void hopper_na2d_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::optional& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& q_shape, + const std::tuple& kv_shape, + const std::tuple& qkv_shape, + const std::tuple& query_tile_shape, + const std::tuple& key_tile_shape, + int kernel_type) { + TORCH_CHECK(query.dim() == 4, "Tensors must be 4-D."); + + hopper_fna_generic_forward( + out, + query, + key, + value, + logsumexp, + kernel_size, + stride, + dilation, + is_causal, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + query_tile_shape, + key_tile_shape, + kernel_type); +} + +void hopper_na3d_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::optional& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& q_shape, + const std::tuple& kv_shape, + const std::tuple& qkv_shape, + const std::tuple& query_tile_shape, + const std::tuple& key_tile_shape, + int kernel_type) { + TORCH_CHECK(query.dim() == 4, "Tensors must be 4-D."); + + hopper_fna_generic_forward( + out, + query, + key, + value, + logsumexp, + kernel_size, + stride, + dilation, + is_causal, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + query_tile_shape, + key_tile_shape, + kernel_type); +} + +} // namespace natten diff --git a/natten/csrc/src/reference_backward.cu b/natten/csrc/src/reference_backward.cu new file mode 100644 index 00000000..a83b6f6f --- /dev/null +++ b/natten/csrc/src/reference_backward.cu @@ -0,0 +1,335 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ +/*! \file + \brief Reference backward +*/ + +#include +#include +#include +#include +#include + +#include +#include + +#ifdef NATTEN_WITH_CUTLASS +#include +#include + +template +auto std_tuple_to_cute_tuple(StdTuple a) { + static_assert( + std::tuple_size_v > 0 && std::tuple_size_v < 4); + + if constexpr (std::tuple_size_v == 1) { + return cute::make_tuple(std::get<0>(a)); + } else if constexpr (std::tuple_size_v == 2) { + return cute::make_tuple(std::get<0>(a), std::get<1>(a)); + } else { + return cute::make_tuple(std::get<0>(a), std::get<1>(a), std::get<2>(a)); + } +} + +#endif + +namespace natten { + +template +void reference_na_generic_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + const StdNADim& kernel_size, + const StdNADim& stride, + const StdNADim& dilation, + const StdCausal& is_causal, + float attn_scale, + const StdNADim& qkv_shape, + int num_extra_kv) { + static_assert( + std::tuple_size_v > 0 && std::tuple_size_v < 4); + static constexpr int kNADim = std::tuple_size_v; + static_assert(std::tuple_size_v == kNADim); + +#ifdef NATTEN_WITH_CUTLASS + at::cuda::OptionalCUDAGuard device_guard(query.device()); + + // TODO: please please simplify these checks!!! + CHECK_CUDA(query); + CHECK_CUDA(key); + CHECK_CUDA(value); + CHECK_CUDA(out); + CHECK_CUDA(grad_query); + CHECK_CUDA(grad_key); + CHECK_CUDA(grad_value); + CHECK_CUDA(grad_out); + CHECK_CUDA(logsumexp); + + CHECK_CONTIGUOUS(query); + CHECK_CONTIGUOUS(key); + CHECK_CONTIGUOUS(value); + CHECK_CONTIGUOUS(grad_query); + CHECK_CONTIGUOUS(grad_key); + CHECK_CONTIGUOUS(grad_value); + CHECK_CONTIGUOUS(out); + CHECK_CONTIGUOUS(grad_out); + CHECK_CONTIGUOUS(logsumexp); + + CheckArgs(kernel_size, stride, dilation); + + CheckIfPropertiesMatch(query, key, value); + CheckIfPropertiesMatch(grad_value, grad_out, out); + CheckIfPropertiesMatch(grad_query, grad_key, grad_value); + CheckIfPropertiesMatch(grad_query, query, value); + + // Everything's flattened to 1D, because we concat additional kvs + CheckIfTensorShapesMatchExceptHeadDim<1>(query, out); + CheckIfTensorShapesMatchExceptHeadDim<1>(key, value); + // CheckIfBatchHeadsHeadDimMatch(query, key); + CheckIfTensorShapesMatch<1>(grad_query, query); + CheckIfTensorShapesMatch<1>(grad_key, key); + CheckIfTensorShapesMatch<1>(grad_value, value); + CheckIfTensorShapesMatch<1>(grad_out, out); + CheckLogSumExp<1>(out, logsumexp); + + TORCH_CHECK( + query.size(0) == key.size(0), + "Blackwell FMHA forward: Query and key must match in batch size, got ", + "query.shape[0]=", + query.size(0), + ", key.shape[0]=", + key.size(0)); + + TORCH_CHECK( + query.size(3) == key.size(3), + "Blackwell FMHA forward: Query and key must match in head dim, got ", + "query.shape[3]=", + query.size(3), + ", key.shape[3]=", + key.size(3)); + + // GQA/MQA is supported + TORCH_CHECK( + query.size(2) >= key.size(2), + "Blackwell FMHA forward: Query heads must be greater than or equal to key/value heads, got ", + "query.shape[2]=", + query.size(2), + ", key.shape[2]=", + key.size(2)); + + TORCH_CHECK( + query.size(2) % key.size(2) == 0, + "Blackwell FMHA forward: Query heads must evenly divide key/value heads, got ", + "query.shape[2]=", + query.size(2), + ", key.shape[2]=", + key.size(2)); + + int batch_size = query.size(0); + int heads = query.size(2); + int heads_kv = key.size(2); + int dim = query.size(3); + + int seqlen_q = query.size(1); + int seqlen_kv = key.size(1); + int dim_value = value.size(3); + + CheckArgsAgainstDim(qkv_shape, kernel_size, dilation); + + auto qkv_shape_ = std_tuple_to_cute_tuple(qkv_shape); + auto window_size = std_tuple_to_cute_tuple(kernel_size); + auto stride_ = std_tuple_to_cute_tuple(stride); + auto dilation_ = std_tuple_to_cute_tuple(dilation); + auto is_causal_ = std_tuple_to_cute_tuple(is_causal); + + int seqlen = size(qkv_shape_); + + TORCH_CHECK( + query.scalar_type() == torch::kFloat16 || + query.scalar_type() == torch::kBFloat16 || + query.scalar_type() == torch::kFloat32, + "Only FP32, FP16, and BF16 operands are supported for now."); + + TORCH_CHECK( + size(qkv_shape_) == seqlen_q, + "Q's sequence length (q.shape[1]) must match the size of QKV shape."); + TORCH_CHECK( + size(qkv_shape_) + num_extra_kv == seqlen_kv, + "KV's sequence length ({k,v}.shape[1]) must match the size of QKV shape + num_extra_kv."); + + int device_id = query.device().index(); + auto cuda_stream = at::cuda::getCurrentCUDAStream(device_id); + + DISPATCH_REFERENCE_FNA_BACKWARD( + kNADim, + query.scalar_type(), + is_causal_, + static_cast(query.data_ptr()), + static_cast(key.data_ptr()), + static_cast(value.data_ptr()), + static_cast(out.data_ptr()), + static_cast(grad_out.data_ptr()), + static_cast(grad_query.data_ptr()), + static_cast(grad_key.data_ptr()), + static_cast(grad_value.data_ptr()), + static_cast(logsumexp.data_ptr()), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_extra_kv, + qkv_shape_, + window_size, + stride_, + dilation_, + attn_scale, + cuda_stream); +#else + TORCH_CHECK(false, "libnatten not compiled with CUTLASS."); +#endif +} + +void reference_na1d_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& qkv_shape, + int num_extra_kv) { + TORCH_CHECK(query.dim() == 4, "Tensors must be 4-D."); + + reference_na_generic_backward( + grad_query, + grad_key, + grad_value, + query, + key, + value, + out, + grad_out, + logsumexp, + kernel_size, + stride, + dilation, + is_causal, + attn_scale, + qkv_shape, + num_extra_kv); +} + +void reference_na2d_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& qkv_shape, + int num_extra_kv) { + TORCH_CHECK(query.dim() == 4, "Tensors must be 4-D."); + + reference_na_generic_backward( + grad_query, + grad_key, + grad_value, + query, + key, + value, + out, + grad_out, + logsumexp, + kernel_size, + stride, + dilation, + is_causal, + attn_scale, + qkv_shape, + num_extra_kv); +} + +void reference_na3d_backward( + at::Tensor& grad_query, + at::Tensor& grad_key, + at::Tensor& grad_value, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + const at::Tensor& out, + const at::Tensor& grad_out, + const at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& qkv_shape, + int num_extra_kv) { + TORCH_CHECK(query.dim() == 4, "Tensors must be 4-D."); + + reference_na_generic_backward( + grad_query, + grad_key, + grad_value, + query, + key, + value, + out, + grad_out, + logsumexp, + kernel_size, + stride, + dilation, + is_causal, + attn_scale, + qkv_shape, + num_extra_kv); +} + +} // namespace natten diff --git a/natten/csrc/src/reference_forward.cu b/natten/csrc/src/reference_forward.cu new file mode 100644 index 00000000..9592ef15 --- /dev/null +++ b/natten/csrc/src/reference_forward.cu @@ -0,0 +1,286 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ +/*! \file + \brief Reference forward +*/ + +#include +#include +#include +#include +#include + +#include +#include + +#ifdef NATTEN_WITH_CUTLASS +#include +#include + +template +auto std_tuple_to_cute_tuple(StdTuple a) { + static_assert( + std::tuple_size_v > 0 && std::tuple_size_v < 4); + + if constexpr (std::tuple_size_v == 1) { + return cute::make_tuple(std::get<0>(a)); + } else if constexpr (std::tuple_size_v == 2) { + return cute::make_tuple(std::get<0>(a), std::get<1>(a)); + } else { + return cute::make_tuple(std::get<0>(a), std::get<1>(a), std::get<2>(a)); + } +} + +#endif + +namespace natten { + +template +void reference_na_generic_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + at::Tensor& logsumexp, + const StdNADim& kernel_size, + const StdNADim& stride, + const StdNADim& dilation, + const StdCausal& is_causal, + float attn_scale, + const StdNADim& qkv_shape, + int num_extra_kv) { + static_assert( + std::tuple_size_v > 0 && std::tuple_size_v < 4); + static constexpr int kNADim = std::tuple_size_v; + static_assert(std::tuple_size_v == kNADim); + +#ifdef NATTEN_WITH_CUTLASS + CHECK_CONTIGUOUS(query); + CHECK_CONTIGUOUS(key); + CHECK_CONTIGUOUS(value); + CHECK_CONTIGUOUS(out); + CHECK_CONTIGUOUS(logsumexp); + + CHECK_CUDA(query); + CHECK_CUDA(key); + CHECK_CUDA(value); + CHECK_CUDA(out); + CHECK_CUDA(logsumexp); + + at::cuda::OptionalCUDAGuard device_guard(query.device()); + + CheckArgs(kernel_size, stride, dilation); + CheckIfPropertiesMatch(query, key, value); + + // Everything's flattened to 1D, because we concat additional kvs + CheckIfTensorShapesMatchExceptHeadDim<1>(query, out); + CheckIfTensorShapesMatchExceptHeadDim<1>(key, value); + // CheckIfBatchHeadsHeadDimMatch(query, key); + CheckLogSumExp<1>(out, logsumexp); + + TORCH_CHECK( + query.size(0) == key.size(0), + "Blackwell FMHA forward: Query and key must match in batch size, got ", + "query.shape[0]=", + query.size(0), + ", key.shape[0]=", + key.size(0)); + + TORCH_CHECK( + query.size(3) == key.size(3), + "Blackwell FMHA forward: Query and key must match in head dim, got ", + "query.shape[3]=", + query.size(3), + ", key.shape[3]=", + key.size(3)); + + // GQA/MQA is supported + TORCH_CHECK( + query.size(2) >= key.size(2), + "Blackwell FMHA forward: Query heads must be greater than or equal to key/value heads, got ", + "query.shape[2]=", + query.size(2), + ", key.shape[2]=", + key.size(2)); + + TORCH_CHECK( + query.size(2) % key.size(2) == 0, + "Blackwell FMHA forward: Query heads must evenly divide key/value heads, got ", + "query.shape[2]=", + query.size(2), + ", key.shape[2]=", + key.size(2)); + + int batch_size = query.size(0); + int heads = query.size(2); + int heads_kv = key.size(2); + int dim = query.size(3); + + int seqlen_q = query.size(1); + int seqlen_kv = key.size(1); + int dim_value = value.size(3); + + CheckArgsAgainstDim(qkv_shape, kernel_size, dilation); + + auto qkv_shape_ = std_tuple_to_cute_tuple(qkv_shape); + auto window_size = std_tuple_to_cute_tuple(kernel_size); + auto stride_ = std_tuple_to_cute_tuple(stride); + auto dilation_ = std_tuple_to_cute_tuple(dilation); + auto is_causal_ = std_tuple_to_cute_tuple(is_causal); + + int seqlen = size(qkv_shape_); + + TORCH_CHECK( + size(qkv_shape_) == seqlen_q, + "Q's sequence length (q.shape[1]) must match the size of QKV shape."); + TORCH_CHECK( + size(qkv_shape_) + num_extra_kv == seqlen_kv, + "KV's sequence length ({k,v}.shape[1]) must match the size of QKV shape + num_extra_kv."); + + TORCH_CHECK( + query.scalar_type() == torch::kFloat16 || + query.scalar_type() == torch::kBFloat16 || + query.scalar_type() == torch::kFloat32, + "Only FP32, FP16, and BF16 operands are supported for now."); + + int device_id = query.device().index(); + auto cuda_stream = at::cuda::getCurrentCUDAStream(device_id); + + DISPATCH_REFERENCE_FNA_FORWARD( + kNADim, + query.scalar_type(), + is_causal_, + static_cast(query.data_ptr()), + static_cast(key.data_ptr()), + static_cast(value.data_ptr()), + static_cast(out.data_ptr()), + static_cast(logsumexp.data_ptr()), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_extra_kv, + qkv_shape_, + window_size, + stride_, + dilation_, + attn_scale, + cuda_stream); +#else + TORCH_CHECK(false, "libnatten not compiled with CUTLASS."); +#endif +} + +void reference_na1d_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& qkv_shape, + int num_extra_kv) { + TORCH_CHECK(query.dim() == 4, "Tensors must be 4-D."); + + reference_na_generic_forward( + out, + query, + key, + value, + logsumexp, + kernel_size, + stride, + dilation, + is_causal, + attn_scale, + qkv_shape, + num_extra_kv); +} + +void reference_na2d_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& qkv_shape, + int num_extra_kv) { + TORCH_CHECK(query.dim() == 4, "Tensors must be 4-D."); + + reference_na_generic_forward( + out, + query, + key, + value, + logsumexp, + kernel_size, + stride, + dilation, + is_causal, + attn_scale, + qkv_shape, + num_extra_kv); +} + +void reference_na3d_forward( + at::Tensor& out, + const at::Tensor& query, + const at::Tensor& key, + const at::Tensor& value, + at::Tensor& logsumexp, + const std::tuple& kernel_size, + const std::tuple& stride, + const std::tuple& dilation, + const std::tuple& is_causal, + float attn_scale, + const std::tuple& qkv_shape, + int num_extra_kv) { + TORCH_CHECK(query.dim() == 4, "Tensors must be 4-D."); + + reference_na_generic_forward( + out, + query, + key, + value, + logsumexp, + kernel_size, + stride, + dilation, + is_causal, + attn_scale, + qkv_shape, + num_extra_kv); +} + +} // namespace natten diff --git a/natten/csrc/src/token_permute.cu b/natten/csrc/src/token_permute.cu new file mode 100644 index 00000000..9415e1e6 --- /dev/null +++ b/natten/csrc/src/token_permute.cu @@ -0,0 +1,268 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ +/*! \file + \brief Token Permute +*/ + +#include +#include +#include +#include +#include + +#include +#include + +#ifdef NATTEN_WITH_CUTLASS +#include + +template +auto std_tuple_to_cute_tuple(StdTuple a) { + static_assert( + std::tuple_size_v > 0 && std::tuple_size_v < 4); + + if constexpr (std::tuple_size_v == 1) { + return cute::make_tuple(static_cast(std::get<0>(a))); + } else if constexpr (std::tuple_size_v == 2) { + return cute::make_tuple( + static_cast(std::get<0>(a)), static_cast(std::get<1>(a))); + } else { + return cute::make_tuple( + static_cast(std::get<0>(a)), + static_cast(std::get<1>(a)), + static_cast(std::get<2>(a))); + } +} + +#endif + +namespace natten { + +template +void token_permute_generic( + at::Tensor& out, + const at::Tensor& in, + const StdNADim& token_layout, + const StdNADim& tile_shape, + const StdNADim& dilation, + bool flip_tiled_dims) { + static_assert( + std::tuple_size_v > 0 && std::tuple_size_v < 4); + static constexpr int kNADim = std::tuple_size_v; + +#ifdef NATTEN_WITH_CUTLASS + CHECK_CONTIGUOUS(out); + CHECK_CONTIGUOUS(in); + CHECK_CUDA(out); + CHECK_CUDA(in); + + CheckIfPropertiesMatch(in, out); + + at::cuda::OptionalCUDAGuard device_guard(in.device()); + + TORCH_CHECK( + in.dim() == 3 + kNADim, + "Token permute ", + kNADim, + "-D expected rank-", + 3 + kNADim, + " tensor input, got ", + "input.dim()=", + in.dim()); + + TORCH_CHECK( + in.size(0) == out.size(0), + "Token permute input and output batches must match, got ", + "input.shape[0]=", + in.size(0), + ", output.shape[0]=", + out.size(0)); + + TORCH_CHECK( + in.size(1 + kNADim) == out.size(2), + "Token permute input and output heads must match, got ", + "input.shape[-2]=", + in.size(1 + kNADim), + ", output.shape[-2]=", + out.size(2)); + + TORCH_CHECK( + in.size(2 + kNADim) == out.size(3), + "Token permute input and output dims must match, got ", + "input.shape[-1]=", + in.size(2 + kNADim), + ", output.shape[-2]=", + out.size(3)); + + int batch = out.size(0); + int seqlen = out.size(1); + int heads = out.size(2); + int dim = out.size(3); + + auto token_layout_ = std_tuple_to_cute_tuple(token_layout); + auto tile_shape_ = std_tuple_to_cute_tuple(tile_shape); + auto dilation_ = std_tuple_to_cute_tuple(dilation); + + TORCH_CHECK( + size(token_layout_) <= seqlen, + "Token permute input's flattened token layout must be <= output sequence size."); + + int device_id = in.device().index(); + auto cuda_stream = at::cuda::getCurrentCUDAStream(device_id); + + cudaDeviceProp* device_props = at::cuda::getDeviceProperties(device_id); + const int cc = device_props->major * 10 + device_props->minor; + + bool is_fp8_allowed = cc == 100 || cc == 103; + + TORCH_CHECK( + in.scalar_type() == torch::kFloat16 || + in.scalar_type() == torch::kBFloat16 || + in.scalar_type() == torch::kFloat32 || + (is_fp8_allowed && + (in.scalar_type() == c10::ScalarType::Float8_e4m3fn || + in.scalar_type() == c10::ScalarType::Float8_e5m2)), + "Only FP32, FP16, BF16, and FP8 operands are supported for now."); + + bool success = false; + if (in.scalar_type() == torch::kFloat32) { + success = natten::tokperm::token_permute_op( + reinterpret_cast(in.data_ptr()), + reinterpret_cast(out.data_ptr()), + batch, + heads, + seqlen, + dim, + token_layout_, + tile_shape_, + dilation_, + flip_tiled_dims, + cuda_stream); + } else if (in.scalar_type() == torch::kFloat16) { + success = natten::tokperm::token_permute_op( + reinterpret_cast(in.data_ptr()), + reinterpret_cast(out.data_ptr()), + batch, + heads, + seqlen, + dim, + token_layout_, + tile_shape_, + dilation_, + flip_tiled_dims, + cuda_stream); + } else if (in.scalar_type() == torch::kBFloat16) { + success = natten::tokperm::token_permute_op( + reinterpret_cast(in.data_ptr()), + reinterpret_cast(out.data_ptr()), + batch, + heads, + seqlen, + dim, + token_layout_, + tile_shape_, + dilation_, + flip_tiled_dims, + cuda_stream); + } +#if defined(NATTEN_WITH_BLACKWELL_FNA) + else if ( + is_fp8_allowed && in.scalar_type() == c10::ScalarType::Float8_e4m3fn) { + success = natten::tokperm::token_permute_op( + reinterpret_cast(in.data_ptr()), + reinterpret_cast(out.data_ptr()), + batch, + heads, + seqlen, + dim, + token_layout_, + tile_shape_, + dilation_, + flip_tiled_dims, + cuda_stream); + } else if ( + is_fp8_allowed && in.scalar_type() == c10::ScalarType::Float8_e5m2) { + success = natten::tokperm::token_permute_op( + reinterpret_cast(in.data_ptr()), + reinterpret_cast(out.data_ptr()), + batch, + heads, + seqlen, + dim, + token_layout_, + tile_shape_, + dilation_, + flip_tiled_dims, + cuda_stream); + } +#endif + + TORCH_CHECK(success, "Token Permute kernel launch failed."); + +#else + TORCH_CHECK(false, "libnatten not compiled with CUTLASS."); +#endif +} + +void token_permute_1d( + at::Tensor& out, + const at::Tensor& in, + const std::tuple& tile_shape, + const std::tuple& dilation, + bool flip_tiled_dims) { + TORCH_CHECK(in.dim() == 4, "Input tensor must be 4-D."); + + token_permute_generic( + out, in, {in.size(1)}, tile_shape, dilation, flip_tiled_dims); +} + +void token_permute_2d( + at::Tensor& out, + const at::Tensor& in, + const std::tuple& tile_shape, + const std::tuple& dilation, + bool flip_tiled_dims) { + TORCH_CHECK(in.dim() == 5, "Input tensor must be 5-D."); + + token_permute_generic( + out, in, {in.size(1), in.size(2)}, tile_shape, dilation, flip_tiled_dims); +} + +void token_permute_3d( + at::Tensor& out, + const at::Tensor& in, + const std::tuple& tile_shape, + const std::tuple& dilation, + bool flip_tiled_dims) { + TORCH_CHECK(in.dim() == 6, "Input tensor must be 6-D."); + + token_permute_generic( + out, + in, + {in.size(1), in.size(2), in.size(3)}, + tile_shape, + dilation, + flip_tiled_dims); +} + +} // namespace natten diff --git a/natten/csrc/src/token_unpermute.cu b/natten/csrc/src/token_unpermute.cu new file mode 100644 index 00000000..09824101 --- /dev/null +++ b/natten/csrc/src/token_unpermute.cu @@ -0,0 +1,268 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ +/*! \file + \brief Token Unpermute +*/ + +#include +#include +#include +#include +#include + +#include +#include + +#ifdef NATTEN_WITH_CUTLASS +#include + +template +auto std_tuple_to_cute_tuple(StdTuple a) { + static_assert( + std::tuple_size_v > 0 && std::tuple_size_v < 4); + + if constexpr (std::tuple_size_v == 1) { + return cute::make_tuple(static_cast(std::get<0>(a))); + } else if constexpr (std::tuple_size_v == 2) { + return cute::make_tuple( + static_cast(std::get<0>(a)), static_cast(std::get<1>(a))); + } else { + return cute::make_tuple( + static_cast(std::get<0>(a)), + static_cast(std::get<1>(a)), + static_cast(std::get<2>(a))); + } +} + +#endif + +namespace natten { + +template +void token_unpermute_generic( + at::Tensor& out, + const at::Tensor& in, + const StdNADim& token_layout, + const StdNADim& tile_shape, + const StdNADim& dilation, + bool flip_tiled_dims) { + static_assert( + std::tuple_size_v > 0 && std::tuple_size_v < 4); + static constexpr int kNADim = std::tuple_size_v; + +#ifdef NATTEN_WITH_CUTLASS + CHECK_CONTIGUOUS(out); + CHECK_CONTIGUOUS(in); + CHECK_CUDA(out); + CHECK_CUDA(in); + + CheckIfPropertiesMatch(in, out); + + at::cuda::OptionalCUDAGuard device_guard(in.device()); + + TORCH_CHECK( + out.dim() == 3 + kNADim, + "Token unpermute ", + kNADim, + "-D expected rank-", + 3 + kNADim, + " tensor output, got ", + "output.dim()=", + out.dim()); + + TORCH_CHECK( + in.size(0) == out.size(0), + "Token unpermute input and output batches must match, got ", + "input.shape[0]=", + in.size(0), + ", output.shape[0]=", + out.size(0)); + + TORCH_CHECK( + out.size(1 + kNADim) == in.size(2), + "Token unpermute input and output heads must match, got ", + "output.shape[-2]=", + out.size(1 + kNADim), + ", input.shape[-2]=", + in.size(2)); + + TORCH_CHECK( + out.size(2 + kNADim) == in.size(3), + "Token unpermute input and output dims must match, got ", + "output.shape[-1]=", + out.size(2 + kNADim), + ", input.shape[-2]=", + in.size(3)); + + int batch = in.size(0); + int seqlen = in.size(1); + int heads = in.size(2); + int dim = in.size(3); + + auto token_layout_ = std_tuple_to_cute_tuple(token_layout); + auto tile_shape_ = std_tuple_to_cute_tuple(tile_shape); + auto dilation_ = std_tuple_to_cute_tuple(dilation); + + TORCH_CHECK( + size(token_layout_) <= seqlen, + "Token unpermute output's flattened token layout must be <= input sequence size."); + + int device_id = in.device().index(); + auto cuda_stream = at::cuda::getCurrentCUDAStream(device_id); + + cudaDeviceProp* device_props = at::cuda::getDeviceProperties(device_id); + const int cc = device_props->major * 10 + device_props->minor; + + bool is_fp8_allowed = cc == 100 || cc == 103; + + TORCH_CHECK( + in.scalar_type() == torch::kFloat16 || + in.scalar_type() == torch::kBFloat16 || + in.scalar_type() == torch::kFloat32 || + (is_fp8_allowed && + (in.scalar_type() == c10::ScalarType::Float8_e4m3fn || + in.scalar_type() == c10::ScalarType::Float8_e5m2)), + "Only FP32, FP16, BF16, and FP8 operands are supported for now."); + + bool success = false; + if (in.scalar_type() == torch::kFloat32) { + success = natten::tokperm::token_unpermute_op( + reinterpret_cast(in.data_ptr()), + reinterpret_cast(out.data_ptr()), + batch, + heads, + dim, + token_layout_, + tile_shape_, + dilation_, + flip_tiled_dims, + cuda_stream); + } else if (in.scalar_type() == torch::kFloat16) { + success = natten::tokperm::token_unpermute_op( + reinterpret_cast(in.data_ptr()), + reinterpret_cast(out.data_ptr()), + batch, + heads, + dim, + token_layout_, + tile_shape_, + dilation_, + flip_tiled_dims, + cuda_stream); + } else if (in.scalar_type() == torch::kBFloat16) { + success = natten::tokperm::token_unpermute_op( + reinterpret_cast(in.data_ptr()), + reinterpret_cast(out.data_ptr()), + batch, + heads, + dim, + token_layout_, + tile_shape_, + dilation_, + flip_tiled_dims, + cuda_stream); + } +#if defined(NATTEN_WITH_BLACKWELL_FNA) + else if ( + is_fp8_allowed && in.scalar_type() == c10::ScalarType::Float8_e4m3fn) { + success = natten::tokperm::token_unpermute_op( + reinterpret_cast(in.data_ptr()), + reinterpret_cast(out.data_ptr()), + batch, + heads, + dim, + token_layout_, + tile_shape_, + dilation_, + flip_tiled_dims, + cuda_stream); + } else if ( + is_fp8_allowed && in.scalar_type() == c10::ScalarType::Float8_e5m2) { + success = natten::tokperm::token_unpermute_op( + reinterpret_cast(in.data_ptr()), + reinterpret_cast(out.data_ptr()), + batch, + heads, + dim, + token_layout_, + tile_shape_, + dilation_, + flip_tiled_dims, + cuda_stream); + } +#endif + + TORCH_CHECK(success, "Token Unpermute kernel launch failed."); + +#else + TORCH_CHECK(false, "libnatten not compiled with CUTLASS."); +#endif +} + +void token_unpermute_1d( + at::Tensor& out, + const at::Tensor& in, + const std::tuple& tile_shape, + const std::tuple& dilation, + bool flip_tiled_dims) { + TORCH_CHECK(out.dim() == 4, "Output tensor must be 4-D."); + + token_unpermute_generic( + out, in, {out.size(1)}, tile_shape, dilation, flip_tiled_dims); +} + +void token_unpermute_2d( + at::Tensor& out, + const at::Tensor& in, + const std::tuple& tile_shape, + const std::tuple& dilation, + bool flip_tiled_dims) { + TORCH_CHECK(out.dim() == 5, "Output tensor must be 5-D."); + + token_unpermute_generic( + out, + in, + {out.size(1), out.size(2)}, + tile_shape, + dilation, + flip_tiled_dims); +} + +void token_unpermute_3d( + at::Tensor& out, + const at::Tensor& in, + const std::tuple& tile_shape, + const std::tuple& dilation, + bool flip_tiled_dims) { + TORCH_CHECK(out.dim() == 6, "Output tensor must be 6-D."); + + token_unpermute_generic( + out, + in, + {out.size(1), out.size(2), out.size(3)}, + tile_shape, + dilation, + flip_tiled_dims); +} + +} // namespace natten diff --git a/natten/flake.lock b/natten/flake.lock new file mode 100644 index 00000000..d7fdf232 --- /dev/null +++ b/natten/flake.lock @@ -0,0 +1,117 @@ +{ + "nodes": { + "flake-compat": { + "locked": { + "lastModified": 1767039857, + "narHash": "sha256-vNpUSpF5Nuw8xvDLj2KCwwksIbjua2LZCqhV1LNRDns=", + "owner": "edolstra", + "repo": "flake-compat", + "rev": "5edf11c44bc78a0d334f6334cdaf7d60d732daab", + "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" + } + }, + "kernel-builder": { + "inputs": { + "flake-compat": "flake-compat", + "flake-utils": "flake-utils", + "nixpkgs": "nixpkgs", + "rust-overlay": "rust-overlay" + }, + "locked": { + "lastModified": 1783417137, + "narHash": "sha256-/iWe/m5ZACdb/DcXTrvXmbHS2QrzgGcGa7hPKi460tE=", + "owner": "huggingface", + "repo": "kernels", + "rev": "81580bb92577f2f7228661ca2a221fc052375709", + "type": "github" + }, + "original": { + "owner": "huggingface", + "repo": "kernels", + "type": "github" + } + }, + "nixpkgs": { + "locked": { + "lastModified": 1776927958, + "narHash": "sha256-XOzEtft7E0P6TgQViLUOQeGHlEYiQ0+FY24BPEksj6s=", + "owner": "NixOS", + "repo": "nixpkgs", + "rev": "fec2c46cca5bf9767486a290abae51200b656d69", + "type": "github" + }, + "original": { + "owner": "NixOS", + "ref": "nixos-unstable-small", + "repo": "nixpkgs", + "type": "github" + } + }, + "root": { + "inputs": { + "kernel-builder": "kernel-builder" + } + }, + "rust-overlay": { + "inputs": { + "nixpkgs": [ + "kernel-builder", + "nixpkgs" + ] + }, + "locked": { + "lastModified": 1776914043, + "narHash": "sha256-qug5r56yW1qOsjSI99l3Jm15JNT9CvS2otkXNRNtrPI=", + "owner": "oxalica", + "repo": "rust-overlay", + "rev": "2d35c4358d7de3a0e606a6e8b27925d981c01cc3", + "type": "github" + }, + "original": { + "owner": "oxalica", + "repo": "rust-overlay", + "type": "github" + } + }, + "systems": { + "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/natten/flake.nix b/natten/flake.nix new file mode 100644 index 00000000..c218afcc --- /dev/null +++ b/natten/flake.nix @@ -0,0 +1,17 @@ +{ + description = "Flake for NATTEN (Neighborhood Attention) kernels"; + + inputs = { + kernel-builder.url = "github:huggingface/kernels"; + }; + + outputs = + { + self, + kernel-builder, + }: + kernel-builder.lib.genKernelFlakeOutputs { + inherit self; + path = ./.; + }; +} diff --git a/natten/scripts/autogen_blackwell_fmha.py b/natten/scripts/autogen_blackwell_fmha.py new file mode 100644 index 00000000..8dd359b1 --- /dev/null +++ b/natten/scripts/autogen_blackwell_fmha.py @@ -0,0 +1,607 @@ +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# This script is intended to emit fused kernel instantiations into +# a variable number of source files, generate appropriate headers +# and a single dispatcher interface, which will be used by the +# NATTEN API to call the kernels. +# +# NOTE: these scripts are heavily under-documented, and +# overly-repetitive, and will be replaced in future PRs. +# Please use it with caution. + +import argparse +import os +from typing import List, Tuple + +DEFAULT_OUTPUT_DIR = "csrc/" + + +SUPPORTED_GEMM_SHAPES_FORWARD = [ + (256, 128, 32), + (256, 128, 64), + (256, 128, 128), +] + + +KERNEL_DECL_TEMPLATE = """ +void {kernel_name}( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); +""" + + +KERNEL_IMPL_TEMPLATE = """ +void {kernel_name}( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) {{ + + using GemmShape = {GEMMShape}; + using Kernel = natten::cuda::fmha_blackwell::KernelForward< + {dtype}, GemmShape, {is_persistent}, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + {dtype}, GemmShape, {is_persistent}, cutlass::fmha::collective::ResidualMask, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + {dtype}, GemmShape, {is_persistent}, cutlass::fmha::collective::CausalMask, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelForward< + {dtype}, + GemmShape, + /* kIsPersistent= */ {is_persistent}, + cutlass::fmha::collective::ResidualMask, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelForward< + {dtype}, + GemmShape, + /* kIsPersistent= */ {is_persistent}, + cutlass::fmha::collective::CausalMask, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) {{ + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({{bytes}}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }}; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{{}}) == 0; + if (no_mask_required) {{ + Kernel kernel; + launch_kernel(kernel); + }} + else if (is_varlen && is_causal) {{ + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + }} + else if (is_varlen && not is_causal) {{ + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + }} + else if (not is_varlen && is_causal) {{ + KernelWithCausalMask kernel; + launch_kernel(kernel); + }} + else {{ + KernelWithResidualMask kernel; + launch_kernel(kernel); + }} +}} +""" + + +class DataType: + def __init__(self, name, short_name, torch_name, bits): + self.name = name + self.bits = bits + self.short_name = short_name + self.torch_name = torch_name + + +Half = DataType("cutlass::half_t", "float16", "torch::kFloat16", 16) +BFloat = DataType("cutlass::bfloat16_t", "bfloat16", "torch::kBFloat16", 16) +E4M3 = DataType("cutlass::float_e4m3_t", "e4m3", "c10::ScalarType::Float8_e4m3fn", 8) +E5M2 = DataType("cutlass::float_e5m2_t", "e5m2", "c10::ScalarType::Float8_e5m2", 8) + + +def iterable_to_static_cute_tuple(shape_in) -> str: + shape = ", ".join([f"cute::Int<{x}>" for x in shape_in]) + return f"cute::tuple<{shape}>" + + +class BlackwellFmhaInstance: + def __init__( + self, + dtype: DataType, + gemm_shape: Tuple[int, int, int], + is_persistent: bool, + ): + assert len(gemm_shape) == 3 + self.gemm_shape = gemm_shape + self.dtype = dtype + self.max_head_dim = gemm_shape[2] + self.is_persistent = is_persistent + + def get_gemm_shape_cute(self) -> str: + return iterable_to_static_cute_tuple(self.gemm_shape) + + def get_name(self) -> str: + name = "blackwell_fmha" + name += f"_{self.dtype.short_name}" + name += "_" + "x".join([str(x) for x in self.gemm_shape]) + if self.is_persistent: + name += "_persistent" + return name + + def get_decl(self) -> str: + return KERNEL_DECL_TEMPLATE.format( + kernel_name=self.get_name(), + ) + + def get_impl(self) -> str: + return KERNEL_IMPL_TEMPLATE.format( + kernel_name=self.get_name(), + GEMMShape=self.get_gemm_shape_cute(), + is_persistent="true" if self.is_persistent else "false", + dtype=self.dtype.name, + ) + + +def write_combined_source_file(path, filename, headers, kernels): + source_head = [] + source_head += ["#ifdef NATTEN_WITH_CUTLASS\n"] + source_head += ["#ifdef NATTEN_WITH_BLACKWELL_FNA\n"] + + source_head += ["#include \n"] + source_head += ["#include \n"] + + source_head += ["#include \n"] + source_head += ["#include \n"] + source_head += ["#include \n"] + source_head += ["#include \n"] + + source_head += ["#include \n"] + source_head += ["#include \n"] + + source_head += ["#include \n"] + + for header in headers: + source_head += [f"#include <{header}>\n"] + + source_head += ["namespace natten { \n"] + source_head += ["namespace cuda { \n"] + source_head += ["namespace fmha_blackwell { \n\n"] + + source_head = "".join(source_head) + + source_body = [] + for kernel in kernels: + source_body += "\n\n" + kernel.get_impl() + "\n\n" + source_body = "".join(source_body) + + source_foot = "".join( + [ + "} // namespace fmha_blackwell \n", + "} // namespace cuda \n", + "} // namespace natten \n", + "#endif \n", + "#endif \n", + "\n", + ] + ) + filename = f"{path}/{filename}" + with open(filename, "w") as f: + f.write(source_head) + f.write(source_body) + f.write(source_foot) + + +class DTypeDispatcher: + def __init__(self): + self.dtypes = [] + self.name = "DISPATCH_BLACKWELL_FMHA_FORWARD" + + def append(self, dtype: DataType): + self.dtypes.append(dtype) + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += f"#define {self.name}(dtype, dim, q_tile_size, kv_tile_size, persistent, ...) \\\n" + dispatcher_str += " [&] { \\\n" + for i, dtype in enumerate(self.dtypes): + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + dispatcher_str += f"if (dtype == {dtype.torch_name})" + dispatcher_str += " { \\\n" + dispatcher_str += " " + dispatcher_str += f" {self.name}_{dtype.short_name}(dim, q_tile_size, kv_tile_size, persistent, __VA_ARGS__); \\\n" + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += ( + ' throw std::runtime_error("Blackwell FMHA forward kernel dispatch failed! ' + + 'It does not support dtype " + std::string(c10::toString(dtype)) + "."' + + "); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class HeadDimDispatcher: + def __init__(self, dtype: DataType): + self.dtype = dtype + self.name = f"DISPATCH_BLACKWELL_FMHA_FORWARD_{self.dtype.short_name}" + + self.dims: List[int] = [] + + def append(self, dim: int): + self.dims.append(dim) + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += ( + f"#define {self.name}(dim, q_tile_size, kv_tile_size, persistent, ...) \\\n" + ) + dispatcher_str += " [&] { \\\n" + for i, dim in enumerate(sorted(self.dims)): + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + dispatcher_str += f"if (dim <= {dim})" + dispatcher_str += " { \\\n" + dispatcher_str += " " + dispatcher_str += f" {self.name}_headdim{dim}(q_tile_size, kv_tile_size, persistent, __VA_ARGS__); \\\n" + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += ( + " throw std::runtime_error(" + '"Blackwell FMHA forward kernel dispatch failed! ' + 'It does not support head dim "' + + f' + std::to_string(dim) + " for {self.dtype.short_name}."' + + "); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class TileSizeDispatcher: + def __init__( + self, + dtype: DataType, + head_dim: int, + ): + self.dtype = dtype + self.head_dim = head_dim + + self.name = ( + f"DISPATCH_BLACKWELL_FMHA_FORWARD_{self.dtype.short_name}_headdim{head_dim}" + ) + + self.tile_sizes: List = [] + + def append(self, tile_size): + assert len(tile_size) == 2 + assert isinstance(tile_size[0], int) + assert isinstance(tile_size[1], int) + self.tile_sizes.append(tile_size) + + def get_kernel_instance(self, q_tile_size, kv_tile_size, persistent): + gemm_M = q_tile_size + gemm_N = kv_tile_size + gemm_K = self.head_dim + gemm_shape = (gemm_M, gemm_N, gemm_K) + + supported_gemm_shapes = SUPPORTED_GEMM_SHAPES_FORWARD + assert gemm_shape in supported_gemm_shapes + + kernel = BlackwellFmhaInstance( + dtype=self.dtype, + gemm_shape=gemm_shape, + is_persistent=persistent, + ) + + return kernel + + def get_target_name(self, q_tile_size, kv_tile_size, persistent): + kernel = self.get_kernel_instance(q_tile_size, kv_tile_size, persistent) + return kernel.get_name() + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += ( + f"#define {self.name}(q_tile_size, kv_tile_size, persistent, ...) \\\n" + ) + dispatcher_str += " [&] { \\\n" + i = 0 + for q_tile_size, kv_tile_size in self.tile_sizes: + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + i += 1 + dispatcher_str += "if (" + dispatcher_str += f"q_tile_size == {q_tile_size}" + dispatcher_str += " && " + dispatcher_str += "\\\n" + dispatcher_str += f"kv_tile_size == {kv_tile_size}" + dispatcher_str += ")" + dispatcher_str += " { \\\n" + + dispatcher_str += " if (persistent) { \\\n" + dispatcher_str += f" natten::cuda::fmha_blackwell::{self.get_target_name(q_tile_size, kv_tile_size, True)}(__VA_ARGS__); \\\n" + dispatcher_str += " } else { \\\n" + dispatcher_str += f" natten::cuda::fmha_blackwell::{self.get_target_name(q_tile_size, kv_tile_size, False)}(__VA_ARGS__); \\\n" + dispatcher_str += " } \\\n" + + dispatcher_str += "} \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += " " + dispatcher_str += ( + " throw std::runtime_error(" + '"Blackwell FMHA forward kernel dispatch failed! ' + "It got invalid Q tile and KV tile combination " + + f"({self.dtype.short_name}, head_dim {self.head_dim}): " + 'q_tile=" + std::to_string(q_tile_size)' + ' + ", kv_tile=" + std::to_string(kv_tile_size)' + ' + "."' + "); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +def write_header_file(content, path, namespaces, extra_includes=None): + extra_includes = extra_includes or [] + header_head = [ + "#pragma once\n", + "\n\n", + ] + header_head += ["#include \n"] + header_head += ["#include \n"] + header_head += ["#include \n"] + header_head += ["#include \n"] + header_head += ["#ifdef NATTEN_WITH_CUTLASS\n"] + header_head += ["#ifdef NATTEN_WITH_BLACKWELL_FNA\n"] + for incl in extra_includes: + header_head += [f"#include <{incl}> \n"] + + for namespace in namespaces: + header_head += [f"namespace {namespace}", " { \n"] + + header_foot = [ + "\n\n", + ] + for namespace in namespaces: + header_foot += ["} ", f"// namespace {namespace}", " \n"] + header_foot += [ + "#endif \n", + "#endif \n", + "\n", + ] + with open(path, "w") as f: + f.write("".join(header_head)) + f.write(content) + f.write("".join(header_foot)) + + +def generate_blackwell_fmha_kernels(path, num_splits=2): + + SUPPORTED_DTYPES = [ + Half, + BFloat, + E4M3, + E5M2, + ] + + HEAD_DIMS = [32, 64, 128] + + TILE_SIZES = [ + (256, 128), + ] + + head_dim_dispatchers = [] + tile_size_dispatchers = [] + kernels = [] + + dtype_dispatcher = DTypeDispatcher() + for dtype in SUPPORTED_DTYPES: + dtype_dispatcher.append(dtype) + + head_dim_dispatcher = HeadDimDispatcher(dtype=dtype) + + for head_dim in HEAD_DIMS: + head_dim_dispatcher.append(head_dim) + + tile_size_dispatcher = TileSizeDispatcher( + dtype=dtype, + head_dim=head_dim, + ) + + for tile_size in TILE_SIZES: + tile_size_dispatcher.append(tile_size) + for persistent in [False, True]: + kernels.append( + tile_size_dispatcher.get_kernel_instance( + tile_size[0], tile_size[1], persistent + ) + ) + tile_size_dispatchers.append(tile_size_dispatcher) + head_dim_dispatchers.append(head_dim_dispatcher) + + # + + path_to_sources = f"{path}/autogen/src/cuda/blackwell_fmha/" + rel_header = "natten_autogen/cuda/blackwell_fmha/" + path_to_header_dir = f"{path}/autogen/include/{rel_header}" + + os.makedirs(path_to_sources, exist_ok=False) + os.makedirs(path_to_header_dir, exist_ok=False) + + path_headers = f"{path_to_header_dir}kernels.h" + path_dtype = f"{path_to_header_dir}interface.h" + path_head_dim = f"{path_to_header_dir}dispatch_head_dim.h" + path_tile_size = f"{path_to_header_dir}dispatch_tile_size.h" + + rel_path_headers = f"{rel_header}kernels.h" + rel_path_head_dim = f"{rel_header}dispatch_head_dim.h" + rel_path_tile_size = f"{rel_header}dispatch_tile_size.h" + + dtype_disp = dtype_dispatcher.get_dispatcher() + + head_dim_disp = "" + for dispatcher in head_dim_dispatchers: + head_dim_disp += dispatcher.get_dispatcher() + + tile_size_disp = "" + for dispatcher in tile_size_dispatchers: + tile_size_disp += dispatcher.get_dispatcher() + + headers = "" + for kernel in kernels: + headers += kernel.get_decl() + + assert ( + len(kernels) >= num_splits + ), f"Generated {len(kernels)} kernels, but got {num_splits=}." + split_size = len(kernels) // num_splits + num_splits_with_res = len(kernels) % num_splits + kernels_emitted = [] + kernels_split = [] + for split_idx in range(num_splits): + kernel_start_idx = split_size * split_idx + min(num_splits_with_res, split_idx) + num_kernels_in_split = split_size + ( + 1 if split_idx < num_splits_with_res else 0 + ) + kernel_end_idx = kernel_start_idx + num_kernels_in_split + assert kernel_end_idx <= len(kernels) + pth_set = set() + source_list = [] + for kernel_idx in range(kernel_start_idx, kernel_end_idx): + kernel = kernels[kernel_idx] + # pth_set.add(kernel.path_to_header) + source_list.append(kernel) + kernels_emitted.append(kernel_idx) + pth_set.add(rel_path_headers) + write_combined_source_file( + path_to_sources, f"source_{split_idx}.cu", sorted(pth_set), source_list + ) + kernels_split.append(source_list) + # print(f"{split_idx=}, {kernel_start_idx=}, {kernel_end_idx=}, {len(kernels_emitted)=}") + assert split_idx == num_splits - 1, f"Expected {split_idx=} == {num_splits=} - 1" + assert len(kernels_emitted) == len(kernels) + assert sorted(kernels_emitted) == [ + x for x in range(len(kernels)) + ], f"{sorted(kernels_emitted)=}" + assert all(len(x) > 0 for x in kernels_split) + + namespaces = ["natten", "cuda", "fmha_blackwell"] + cuda_headers = [ + "natten/natten.h", + "ATen/ATen.h", + "ATen/cuda/CUDAContext.h", + "c10/cuda/CUDAGuard.h", + "c10/cuda/CUDAStream.h", + "natten/natten.h", + "natten/helpers.h", + "natten/cuda/fmha_blackwell/fmha_forward.cuh", + ] + write_header_file( + dtype_disp, path_dtype, namespaces, cuda_headers + [rel_path_head_dim] + ) + write_header_file( + head_dim_disp, path_head_dim, namespaces, cuda_headers + [rel_path_tile_size] + ) + write_header_file( + tile_size_disp, path_tile_size, namespaces, cuda_headers + [rel_path_headers] + ) + write_header_file(headers, path_headers, namespaces, cuda_headers) + + +def generate_blackwell_fmha(output_directory: str, num_splits: int): + generate_blackwell_fmha_kernels(output_directory, num_splits=num_splits) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "-o", + "--output-directory", + default=DEFAULT_OUTPUT_DIR, + help="Path to the directory where the auto-generated " + "kernel instantiations are dumped. " + f"Default: {DEFAULT_OUTPUT_DIR}", + ) + parser.add_argument( + "--num-splits", + type=int, + default=1, + help="Number of source files into which the kernels are split. Default: 1.", + ) + args = parser.parse_args() + generate_blackwell_fmha(args.output_directory, args.num_splits) diff --git a/natten/scripts/autogen_blackwell_fmha_bwd.py b/natten/scripts/autogen_blackwell_fmha_bwd.py new file mode 100644 index 00000000..9a6ba0d3 --- /dev/null +++ b/natten/scripts/autogen_blackwell_fmha_bwd.py @@ -0,0 +1,623 @@ +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# This script is intended to emit fused kernel instantiations into +# a variable number of source files, generate appropriate headers +# and a single dispatcher interface, which will be used by the +# NATTEN API to call the kernels. +# +# NOTE: these scripts are heavily under-documented, and +# overly-repetitive, and will be replaced in future PRs. +# Please use it with caution. + +import argparse +import os +from typing import List, Tuple + +DEFAULT_OUTPUT_DIR = "csrc/" + +SUPPORTED_GEMM_SHAPES_BACKWARD = [ + (128, 128, 32), + (128, 128, 64), + (128, 128, 128), +] + + +KERNEL_DECL_TEMPLATE = """ +void {kernel_name}( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + bool deterministic, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); +""" + + +KERNEL_IMPL_TEMPLATE = """ +void {kernel_name}( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + bool deterministic, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) {{ + + auto run = [&](auto kIsDeterministic) {{ + constexpr bool IsDeterministic = decltype(kIsDeterministic)::value; + + using GemmShape = {GEMMShape}; + using Kernel = natten::cuda::fmha_blackwell::KernelBackward< + {dtype}, GemmShape, cutlass::fmha::collective::NoMask, /* kIsVarlen= */ false, IsDeterministic>; + using KernelWithResidualMask = natten::cuda::fmha_blackwell::KernelBackward< + {dtype}, GemmShape, cutlass::fmha::collective::ResidualMaskForBackward, /* kIsVarlen= */ false, IsDeterministic>; + using KernelWithCausalMask = natten::cuda::fmha_blackwell::KernelBackward< + {dtype}, GemmShape, cutlass::fmha::collective::CausalForBackwardMask, /* kIsVarlen= */ false, IsDeterministic>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_blackwell::KernelBackward< + {dtype}, + GemmShape, + cutlass::fmha::collective::ResidualMaskForBackward, + /* kIsVarlen= */ true, IsDeterministic>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_blackwell::KernelBackward< + {dtype}, + GemmShape, + cutlass::fmha::collective::CausalForBackwardMask, + /* kIsVarlen= */ true, IsDeterministic>; + + int* dq_semaphore_ptr = nullptr; + at::Tensor dq_semaphore; + if constexpr (IsDeterministic) {{ + auto kBlockM = cute::get<0>(GemmShape{{}}); + int effective_seqlen_q = is_varlen ? max_seqlen_Q : seqlen_q; + dq_semaphore = torch::zeros( + {{(effective_seqlen_q + kBlockM - 1) / kBlockM, batch_size, heads_q}}, + tensor_options.dtype(at::ScalarType::Int)); + dq_semaphore_ptr = static_cast(dq_semaphore.data_ptr()); + }} + + auto launch_kernel = [&](auto& kernel) {{ + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + dq_semaphore_ptr, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({{bytes}}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }}; + + bool no_mask_required = not is_varlen && not is_causal && + seqlen_q % get<0>(GemmShape{{}}) == 0 && seqlen_k % get<1>(GemmShape{{}}) == 0; + + if (no_mask_required) {{ + Kernel kernel; + launch_kernel(kernel); + }} + else if (is_varlen && is_causal) {{ + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + }} + else if (is_varlen && not is_causal) {{ + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + }} + else if (not is_varlen && is_causal) {{ + KernelWithCausalMask kernel; + launch_kernel(kernel); + }} + else {{ + KernelWithResidualMask kernel; + launch_kernel(kernel); + }} + }}; + + if (deterministic) {{ + run(std::true_type{{}}); + }} else {{ + run(std::false_type{{}}); + }} +}} +""" + + +class DataType: + def __init__(self, name, short_name, torch_name, bits): + self.name = name + self.bits = bits + self.short_name = short_name + self.torch_name = torch_name + + +Half = DataType("cutlass::half_t", "float16", "torch::kFloat16", 16) +BFloat = DataType("cutlass::bfloat16_t", "bfloat16", "torch::kBFloat16", 16) + + +def iterable_to_static_cute_tuple(shape_in) -> str: + shape = ", ".join([f"cute::Int<{x}>" for x in shape_in]) + return f"cute::tuple<{shape}>" + + +class BlackwellFmhaInstance: + def __init__( + self, + dtype: DataType, + gemm_shape: Tuple[int, int, int], + ): + assert len(gemm_shape) == 3 + self.gemm_shape = gemm_shape + self.dtype = dtype + self.max_head_dim = gemm_shape[2] + + def get_gemm_shape_cute(self) -> str: + return iterable_to_static_cute_tuple(self.gemm_shape) + + def get_name(self) -> str: + name = "blackwell_fmha_backward" + name += f"_{self.dtype.short_name}" + name += "_" + "x".join([str(x) for x in self.gemm_shape]) + return name + + def get_decl(self) -> str: + return KERNEL_DECL_TEMPLATE.format( + kernel_name=self.get_name(), + ) + + def get_impl(self) -> str: + return KERNEL_IMPL_TEMPLATE.format( + kernel_name=self.get_name(), + GEMMShape=self.get_gemm_shape_cute(), + dtype=self.dtype.name, + ) + + +def write_combined_source_file(path, filename, headers, kernels): + source_head = [] + source_head += ["#ifdef NATTEN_WITH_CUTLASS\n"] + source_head += ["#ifdef NATTEN_WITH_BLACKWELL_FNA\n"] + + source_head += ["#include \n"] + source_head += ["#include \n"] + + source_head += ["#include \n"] + source_head += ["#include \n"] + source_head += ["#include \n"] + source_head += ["#include \n"] + + source_head += ["#include \n"] + source_head += ["#include \n"] + + source_head += ["#include \n"] + + for header in headers: + source_head += [f"#include <{header}>\n"] + + source_head += ["namespace natten { \n"] + source_head += ["namespace cuda { \n"] + source_head += ["namespace fmha_blackwell { \n\n"] + + source_head = "".join(source_head) + + source_body = [] + for kernel in kernels: + source_body += "\n\n" + kernel.get_impl() + "\n\n" + source_body = "".join(source_body) + + source_foot = "".join( + [ + "} // namespace fmha_blackwell \n", + "} // namespace cuda \n", + "} // namespace natten \n", + "#endif \n", + "#endif \n", + "\n", + ] + ) + filename = f"{path}/{filename}" + with open(filename, "w") as f: + f.write(source_head) + f.write(source_body) + f.write(source_foot) + + +class DTypeDispatcher: + def __init__(self): + self.dtypes = [] + self.name = "DISPATCH_BLACKWELL_FMHA_BACKWARD" + + def append(self, dtype: DataType): + self.dtypes.append(dtype) + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += ( + f"#define {self.name}(dtype, dim, q_tile_size, kv_tile_size, ...) \\\n" + ) + dispatcher_str += " [&] { \\\n" + for i, dtype in enumerate(self.dtypes): + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + dispatcher_str += f"if (dtype == {dtype.torch_name})" + dispatcher_str += " { \\\n" + dispatcher_str += " " + dispatcher_str += f" {self.name}_{dtype.short_name}(dim, q_tile_size, kv_tile_size, __VA_ARGS__); \\\n" + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += ( + ' throw std::runtime_error("Blackwell FMHA backward kernel dispatch failed! ' + + 'It does not support dtype " + std::string(c10::toString(dtype)) + "."' + + "); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class HeadDimDispatcher: + def __init__(self, dtype: DataType): + self.dtype = dtype + self.name = f"DISPATCH_BLACKWELL_FMHA_BACKWARD_{self.dtype.short_name}" + + self.dims: List[int] = [] + + def append(self, dim: int): + self.dims.append(dim) + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += ( + f"#define {self.name}(dim, q_tile_size, kv_tile_size, ...) \\\n" + ) + dispatcher_str += " [&] { \\\n" + for i, dim in enumerate(sorted(self.dims)): + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + dispatcher_str += f"if (dim <= {dim})" + dispatcher_str += " { \\\n" + dispatcher_str += " " + dispatcher_str += f" {self.name}_headdim{dim}(q_tile_size, kv_tile_size, __VA_ARGS__); \\\n" + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += ( + " throw std::runtime_error(" + '"Blackwell FMHA backward kernel dispatch failed! ' + 'It does not support head dim "' + + f' + std::to_string(dim) + " for {self.dtype.short_name}."' + + "); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class TileSizeDispatcher: + def __init__( + self, + dtype: DataType, + head_dim: int, + ): + self.dtype = dtype + self.head_dim = head_dim + + self.name = f"DISPATCH_BLACKWELL_FMHA_BACKWARD_{self.dtype.short_name}_headdim{head_dim}" + + self.tile_sizes: List = [] + + def append(self, tile_size): + assert len(tile_size) == 2 + assert isinstance(tile_size[0], int) + assert isinstance(tile_size[1], int) + self.tile_sizes.append(tile_size) + + def get_kernel_instance(self, q_tile_size, kv_tile_size): + gemm_M = q_tile_size + gemm_N = kv_tile_size + gemm_K = self.head_dim + gemm_shape = (gemm_M, gemm_N, gemm_K) + + supported_gemm_shapes = SUPPORTED_GEMM_SHAPES_BACKWARD + assert gemm_shape in supported_gemm_shapes + + kernel = BlackwellFmhaInstance( + dtype=self.dtype, + gemm_shape=gemm_shape, + ) + + return kernel + + def get_target_name(self, q_tile_size, kv_tile_size): + kernel = self.get_kernel_instance(q_tile_size, kv_tile_size) + return kernel.get_name() + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += f"#define {self.name}(q_tile_size, kv_tile_size, ...) \\\n" + dispatcher_str += " [&] { \\\n" + i = 0 + for q_tile_size, kv_tile_size in self.tile_sizes: + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + i += 1 + dispatcher_str += "if (" + dispatcher_str += f"q_tile_size == {q_tile_size}" + dispatcher_str += " && " + dispatcher_str += "\\\n" + dispatcher_str += f"kv_tile_size == {kv_tile_size}" + dispatcher_str += ")" + dispatcher_str += " { \\\n" + + dispatcher_str += f" natten::cuda::fmha_blackwell::{self.get_target_name(q_tile_size, kv_tile_size)}(__VA_ARGS__); \\\n" + + dispatcher_str += "} \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += " " + dispatcher_str += ( + " throw std::runtime_error(" + '"Blackwell FMHA backward kernel dispatch failed! ' + "It got invalid Q tile and KV tile combination " + + f"({self.dtype.short_name}, head_dim {self.head_dim}): " + 'q_tile=" + std::to_string(q_tile_size)' + ' + ", kv_tile=" + std::to_string(kv_tile_size)' + ' + "."' + "); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +def write_header_file(content, path, namespaces, extra_includes=None): + extra_includes = extra_includes or [] + header_head = [ + "#pragma once\n", + "\n\n", + ] + header_head += ["#include \n"] + header_head += ["#include \n"] + header_head += ["#include \n"] + header_head += ["#include \n"] + header_head += ["#ifdef NATTEN_WITH_CUTLASS\n"] + header_head += ["#ifdef NATTEN_WITH_BLACKWELL_FNA\n"] + for incl in extra_includes: + header_head += [f"#include <{incl}> \n"] + + for namespace in namespaces: + header_head += [f"namespace {namespace}", " { \n"] + + header_foot = [ + "\n\n", + ] + for namespace in namespaces: + header_foot += ["} ", f"// namespace {namespace}", " \n"] + header_foot += [ + "#endif \n", + "#endif \n", + "\n", + ] + with open(path, "w") as f: + f.write("".join(header_head)) + f.write(content) + f.write("".join(header_foot)) + + +def generate_blackwell_fmha_kernels(path, num_splits=2): + + SUPPORTED_DTYPES = [ + Half, + BFloat, + ] + + HEAD_DIMS = [32, 64, 128] + + TILE_SIZES = [ + (128, 128), + ] + + head_dim_dispatchers = [] + tile_size_dispatchers = [] + kernels = [] + + dtype_dispatcher = DTypeDispatcher() + for dtype in SUPPORTED_DTYPES: + dtype_dispatcher.append(dtype) + + head_dim_dispatcher = HeadDimDispatcher(dtype=dtype) + + for head_dim in HEAD_DIMS: + head_dim_dispatcher.append(head_dim) + + tile_size_dispatcher = TileSizeDispatcher( + dtype=dtype, + head_dim=head_dim, + ) + + for tile_size in TILE_SIZES: + tile_size_dispatcher.append(tile_size) + kernels.append( + tile_size_dispatcher.get_kernel_instance(tile_size[0], tile_size[1]) + ) + tile_size_dispatchers.append(tile_size_dispatcher) + head_dim_dispatchers.append(head_dim_dispatcher) + + # + + path_to_sources = f"{path}/autogen/src/cuda/blackwell_fmha_bwd/" + rel_header = "natten_autogen/cuda/blackwell_fmha_bwd/" + path_to_header_dir = f"{path}/autogen/include/{rel_header}" + + os.makedirs(path_to_sources, exist_ok=False) + os.makedirs(path_to_header_dir, exist_ok=False) + + path_headers = f"{path_to_header_dir}kernels.h" + path_dtype = f"{path_to_header_dir}interface.h" + path_head_dim = f"{path_to_header_dir}dispatch_head_dim.h" + path_tile_size = f"{path_to_header_dir}dispatch_tile_size.h" + + rel_path_headers = f"{rel_header}kernels.h" + rel_path_head_dim = f"{rel_header}dispatch_head_dim.h" + rel_path_tile_size = f"{rel_header}dispatch_tile_size.h" + + dtype_disp = dtype_dispatcher.get_dispatcher() + + head_dim_disp = "" + for dispatcher in head_dim_dispatchers: + head_dim_disp += dispatcher.get_dispatcher() + + tile_size_disp = "" + for dispatcher in tile_size_dispatchers: + tile_size_disp += dispatcher.get_dispatcher() + + headers = "" + for kernel in kernels: + headers += kernel.get_decl() + + assert ( + len(kernels) >= num_splits + ), f"Generated {len(kernels)} kernels, but got {num_splits=}." + split_size = len(kernels) // num_splits + num_splits_with_res = len(kernels) % num_splits + kernels_emitted = [] + kernels_split = [] + for split_idx in range(num_splits): + kernel_start_idx = split_size * split_idx + min(num_splits_with_res, split_idx) + num_kernels_in_split = split_size + ( + 1 if split_idx < num_splits_with_res else 0 + ) + kernel_end_idx = kernel_start_idx + num_kernels_in_split + assert kernel_end_idx <= len(kernels) + pth_set = set() + source_list = [] + for kernel_idx in range(kernel_start_idx, kernel_end_idx): + kernel = kernels[kernel_idx] + # pth_set.add(kernel.path_to_header) + source_list.append(kernel) + kernels_emitted.append(kernel_idx) + pth_set.add(rel_path_headers) + write_combined_source_file( + path_to_sources, f"source_{split_idx}.cu", sorted(pth_set), source_list + ) + kernels_split.append(source_list) + # print(f"{split_idx=}, {kernel_start_idx=}, {kernel_end_idx=}, {len(kernels_emitted)=}") + assert split_idx == num_splits - 1, f"Expected {split_idx=} == {num_splits=} - 1" + assert len(kernels_emitted) == len(kernels) + assert sorted(kernels_emitted) == [ + x for x in range(len(kernels)) + ], f"{sorted(kernels_emitted)=}" + assert all(len(x) > 0 for x in kernels_split) + + namespaces = ["natten", "cuda", "fmha_blackwell"] + cuda_headers = [ + "natten/natten.h", + "ATen/ATen.h", + "ATen/cuda/CUDAContext.h", + "c10/cuda/CUDAGuard.h", + "c10/cuda/CUDAStream.h", + "natten/natten.h", + "natten/helpers.h", + "natten/cuda/fmha_blackwell/fmha_backward.cuh", + ] + write_header_file( + dtype_disp, path_dtype, namespaces, cuda_headers + [rel_path_head_dim] + ) + write_header_file( + head_dim_disp, path_head_dim, namespaces, cuda_headers + [rel_path_tile_size] + ) + write_header_file( + tile_size_disp, path_tile_size, namespaces, cuda_headers + [rel_path_headers] + ) + write_header_file(headers, path_headers, namespaces, cuda_headers) + + +def generate_blackwell_fmha(output_directory: str, num_splits: int): + generate_blackwell_fmha_kernels(output_directory, num_splits=num_splits) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "-o", + "--output-directory", + default=DEFAULT_OUTPUT_DIR, + help="Path to the directory where the auto-generated " + "kernel instantiations are dumped. " + f"Default: {DEFAULT_OUTPUT_DIR}", + ) + parser.add_argument( + "--num-splits", + type=int, + default=1, + help="Number of source files into which the kernels are split. Default: 1.", + ) + args = parser.parse_args() + generate_blackwell_fmha(args.output_directory, args.num_splits) diff --git a/natten/scripts/autogen_blackwell_fna.py b/natten/scripts/autogen_blackwell_fna.py new file mode 100644 index 00000000..502fd365 --- /dev/null +++ b/natten/scripts/autogen_blackwell_fna.py @@ -0,0 +1,793 @@ +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# This script is intended to emit fused kernel instantiations into +# a variable number of source files, generate appropriate headers +# and a single dispatcher interface, which will be used by the +# NATTEN API to call the kernels. +# +# NOTE: these scripts are heavily under-documented, and +# overly-repetitive, and will be replaced in future PRs. +# Please use it with caution. + +import argparse +import math +import os +from typing import List, Tuple + +DEFAULT_OUTPUT_DIR = "csrc/" + + +SUPPORTED_GEMM_SHAPES_FORWARD = [ + (256, 128, 32), + (256, 128, 64), + (256, 128, 128), +] + + +KERNEL_DECL_TEMPLATE = """ +void {kernel_name}( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + {DimType} q_shape, + {DimType} kv_shape, + {DimType} qkv_shape, + {DimType} window_size, + {DimType} stride, + {DimType} dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); +""" + + +KERNEL_IMPL_TEMPLATE = """ +void {kernel_name}( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna / fusion parameters + {DimType} q_shape, + {DimType} kv_shape, + {DimType} qkv_shape, + {DimType} window_size, + {DimType} stride, + {DimType} dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) {{ + + using Causal = {Causal}; + using QTileShape = {QTileShape}; + using KVTileShape = {KVTileShape}; + using Config = {GEMMShape}; + using Kernel = natten::cuda::fna_blackwell::KernelForward< + {dtype}, Causal, QTileShape, KVTileShape, Config, {is_persistent}>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({{bytes}}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +}} +""" + + +class DataType: + def __init__(self, name, short_name, torch_name, bits): + self.name = name + self.bits = bits + self.short_name = short_name + self.torch_name = torch_name + + +Half = DataType("cutlass::half_t", "float16", "torch::kFloat16", 16) +BFloat = DataType("cutlass::bfloat16_t", "bfloat16", "torch::kBFloat16", 16) +E4M3 = DataType("cutlass::float_e4m3_t", "e4m3", "c10::ScalarType::Float8_e4m3fn", 8) +E5M2 = DataType("cutlass::float_e5m2_t", "e5m2", "c10::ScalarType::Float8_e5m2", 8) + + +def iterable_to_static_cute_tuple(shape_in) -> str: + shape = ", ".join([f"cute::Int<{x}>" for x in shape_in]) + return f"cute::tuple<{shape}>" + + +def get_dim_type(na_dim: int) -> str: + shape = ", ".join(["int" for _ in range(na_dim)]) + return f"cute::tuple<{shape}>" + + +class BlackwellFnaInstance: + def __init__( + self, + na_dim: int, + dtype: DataType, + gemm_shape: Tuple[int, int, int], + q_tile_shape: tuple, + kv_tile_shape: tuple, + causal: tuple, + is_persistent: bool, + ): + assert 0 < na_dim <= 3 + assert len(gemm_shape) == 3 + assert na_dim == len(q_tile_shape) + assert na_dim == len(kv_tile_shape) + assert na_dim == len(causal) + self.na_dim = na_dim + self.gemm_shape = gemm_shape + self.q_tile_shape = q_tile_shape + self.kv_tile_shape = kv_tile_shape + self.causal = causal + self.dtype = dtype + self.max_head_dim = gemm_shape[2] + self.is_persistent = is_persistent + + def get_q_tile_shape_cute(self) -> str: + return iterable_to_static_cute_tuple(self.q_tile_shape) + + def get_kv_tile_shape_cute(self) -> str: + return iterable_to_static_cute_tuple(self.kv_tile_shape) + + def get_gemm_shape_cute(self) -> str: + return iterable_to_static_cute_tuple(self.gemm_shape) + + def get_causal_cute(self) -> str: + consts = ", ".join( + ["cute::true_type" if c else "cute::false_type" for c in self.causal] + ) + return f"cute::tuple<{consts}>" + + def get_name(self) -> str: + name = f"blackwell_fna{self.na_dim}d" + name += f"_{self.dtype.short_name}" + name += "_" + "x".join([str(x) for x in self.gemm_shape]) + name += "_Q" + "x".join([str(x) for x in self.q_tile_shape]) + name += "_KV" + "x".join([str(x) for x in self.kv_tile_shape]) + name += "_causal" + "x".join(["1" if c else "0" for c in self.causal]) + if self.is_persistent: + name += "_persistent" + return name + + def get_decl(self) -> str: + return KERNEL_DECL_TEMPLATE.format( + kernel_name=self.get_name(), DimType=get_dim_type(self.na_dim) + ) + + def get_impl(self) -> str: + return KERNEL_IMPL_TEMPLATE.format( + kernel_name=self.get_name(), + DimType=get_dim_type(self.na_dim), + Causal=self.get_causal_cute(), + QTileShape=self.get_q_tile_shape_cute(), + KVTileShape=self.get_kv_tile_shape_cute(), + GEMMShape=self.get_gemm_shape_cute(), + is_persistent="true" if self.is_persistent else "false", + dtype=self.dtype.name, + ) + + +def write_combined_source_file(path, filename, headers, kernels): + source_head = [] + source_head += ["#ifdef NATTEN_WITH_CUTLASS\n"] + source_head += ["#ifdef NATTEN_WITH_BLACKWELL_FNA\n"] + + source_head += ["#include \n"] + source_head += ["#include \n"] + + source_head += ["#include \n"] + source_head += ["#include \n"] + source_head += ["#include \n"] + source_head += ["#include \n"] + + source_head += ["#include \n"] + source_head += ["#include \n"] + + source_head += ["#include \n"] + + for header in headers: + source_head += [f"#include <{header}>\n"] + + source_head += ["namespace natten { \n"] + source_head += ["namespace cuda { \n"] + source_head += ["namespace fna_blackwell { \n\n"] + + source_head = "".join(source_head) + + source_body = [] + for kernel in kernels: + source_body += "\n\n" + kernel.get_impl() + "\n\n" + source_body = "".join(source_body) + + source_foot = "".join( + [ + "} // namespace fna_blackwell \n", + "} // namespace cuda \n", + "} // namespace natten \n", + "#endif \n", + "#endif \n", + "\n", + ] + ) + filename = f"{path}/{filename}" + with open(filename, "w") as f: + f.write(source_head) + f.write(source_body) + f.write(source_foot) + + +class NaDimDispatcher: + def __init__(self): + self.name = "DISPATCH_BLACKWELL_FNA_FORWARD" + self.dims = [] + + def append(self, na_dim: int): + self.dims.append(na_dim) + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += f"#define {self.name}(rank, dtype, dim, is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \\\n" + dispatcher_str += " [&] { \\\n" + for i, na_dim in enumerate(self.dims): + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + dispatcher_str += f"if constexpr (rank == {na_dim})" + dispatcher_str += " { \\\n" + dispatcher_str += " " + dispatcher_str += f" {self.name}_{na_dim}D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \\\n" + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += ( + ' throw std::runtime_error("Blackwell FNA forward kernel dispatch failed! ' + + "It only supports NA1D, 2D, and 3D!" + + '"); \\\n' + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class DTypeDispatcher: + def __init__(self, na_dim: int): + self.dtypes: List[DataType] = [] + self.na_dim = na_dim + self.name = f"DISPATCH_BLACKWELL_FNA_FORWARD_{self.na_dim}D" + + def append(self, dtype: DataType): + self.dtypes.append(dtype) + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += f"#define {self.name}(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \\\n" + dispatcher_str += " [&] { \\\n" + for i, dtype in enumerate(self.dtypes): + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + dispatcher_str += f"if (dtype == {dtype.torch_name})" + dispatcher_str += " { \\\n" + dispatcher_str += " " + dispatcher_str += f" {self.name}_{dtype.short_name}(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \\\n" + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += ( + ' throw std::runtime_error("Blackwell FNA-' + + f"{self.na_dim}D forward kernel dispatch failed! " + 'It does not support dtype "' + ' + std::string(c10::toString(dtype)) + "."' + "); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class HeadDimDispatcher: + def __init__(self, na_dim: int, dtype: DataType): + self.na_dim = na_dim + self.dtype = dtype + self.name = ( + f"DISPATCH_BLACKWELL_FNA_FORWARD_{self.na_dim}D_{self.dtype.short_name}" + ) + + self.dims: List[int] = [] + + def append(self, dim: int): + self.dims.append(dim) + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += f"#define {self.name}(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \\\n" + dispatcher_str += " [&] { \\\n" + for i, dim in enumerate(sorted(self.dims)): + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + dispatcher_str += f"if (dim <= {dim})" + dispatcher_str += " { \\\n" + dispatcher_str += " " + dispatcher_str += f" {self.name}_headdim{dim}(is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \\\n" + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += ( + " throw std::runtime_error(" + '"Blackwell FNA forward kernel dispatch failed! ' + 'It does not support head dim "' + + f' + std::to_string(dim) + " for {self.dtype.short_name}."' + + "); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class CausalMaskDispatcher: + def __init__( + self, + na_dim: int, + dtype: DataType, + head_dim: int, + ): + self.na_dim = na_dim + self.dtype = dtype + self.head_dim = head_dim + + self.name = f"DISPATCH_BLACKWELL_FNA_FORWARD_{self.na_dim}D_{self.dtype.short_name}_headdim{head_dim}" + + self.cms: List = [] + + def append(self, cm): + assert len(cm) == self.na_dim, f"{cm} incompatible for {self.na_dim}D NA." + self.cms.append(cm) + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += f"#define {self.name}(is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \\\n" + dispatcher_str += " [&] { \\\n" + i = 0 + for cm in self.cms: + cm_str = "causal" + "x".join(["1" if c else "0" for c in cm]) + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + i += 1 + dispatcher_str += "if (" + for dim in range(self.na_dim): + dispatcher_str += ( + f"cute::get<{dim}>(is_causal)" + if cm[dim] + else f"not cute::get<{dim}>(is_causal)" + ) + if dim != self.na_dim - 1: + dispatcher_str += " && " + dispatcher_str += ")" + dispatcher_str += " { \\\n" + dispatcher_str += " " + dispatcher_str += f" {self.name}_{cm_str}(q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \\\n" + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += " " + dispatcher_str += ( + ' throw std::runtime_error("Blackwell FNA-' + + f"{self.na_dim}D forward kernel dispatch failed! " + f"Causal mask dispatcher ({self.dtype.short_name}, " + f'head_dim {self.head_dim}) got invalid causal mask!"' + "); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class TileShapeDispatcher: + def __init__( + self, + na_dim: int, + dtype: DataType, + head_dim: int, + causal_mask: tuple, + ): + self.na_dim = na_dim + self.dtype = dtype + self.head_dim = head_dim + self.causal_mask = causal_mask + cm_str = "_causal" + "x".join(["1" if c else "0" for c in self.causal_mask]) + + self.name = f"DISPATCH_BLACKWELL_FNA_FORWARD_{self.na_dim}D_{self.dtype.short_name}_headdim{head_dim}" + self.name += cm_str + + self.tile_shapes: List = [] + + def append(self, tile_shape): + assert len(tile_shape) == 2 + assert len(tile_shape[0]) == self.na_dim + assert len(tile_shape[1]) == self.na_dim + self.tile_shapes.append(tile_shape) + + def get_kernel_instance(self, q_tile_shape, kv_tile_shape, persistent): + gemm_M = math.prod(q_tile_shape) + gemm_N = math.prod(kv_tile_shape) + gemm_K = self.head_dim + gemm_shape = (gemm_M, gemm_N, gemm_K) + + assert gemm_shape in SUPPORTED_GEMM_SHAPES_FORWARD + + kernel = BlackwellFnaInstance( + na_dim=self.na_dim, + dtype=self.dtype, + gemm_shape=gemm_shape, + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + causal=self.causal_mask, + is_persistent=persistent, + ) + + return kernel + + def get_target_name(self, q_tile_shape, kv_tile_shape, persistent): + kernel = self.get_kernel_instance(q_tile_shape, kv_tile_shape, persistent) + return kernel.get_name() + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += ( + f"#define {self.name}(q_tile_shape, kv_tile_shape, persistent, ...) \\\n" + ) + dispatcher_str += " [&] { \\\n" + i = 0 + for q_tile_shape, kv_tile_shape in self.tile_shapes: + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + i += 1 + dispatcher_str += "if (" + for dim in range(self.na_dim): + dispatcher_str += ( + f"cute::get<{dim}>(q_tile_shape) == {q_tile_shape[dim]}" + ) + dispatcher_str += " && " + dispatcher_str += "\\\n" + for dim in range(self.na_dim): + dispatcher_str += ( + f"cute::get<{dim}>(kv_tile_shape) == {kv_tile_shape[dim]}" + ) + if dim != self.na_dim - 1: + dispatcher_str += " && " + dispatcher_str += ")" + dispatcher_str += " { \\\n" + + dispatcher_str += " if (persistent) { \\\n" + dispatcher_str += f" natten::cuda::fna_blackwell::{self.get_target_name(q_tile_shape, kv_tile_shape, True)}(__VA_ARGS__); \\\n" + dispatcher_str += " } else { \\\n" + dispatcher_str += f" natten::cuda::fna_blackwell::{self.get_target_name(q_tile_shape, kv_tile_shape, False)}(__VA_ARGS__); \\\n" + dispatcher_str += " } \\\n" + + dispatcher_str += "} \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += " " + q_tile_str = ' + "," + '.join( + [ + f"std::to_string(cute::get<{i}>(q_tile_shape))" + for i in range(self.na_dim) + ] + ) + kv_tile_str = ' + "," + '.join( + [ + f"std::to_string(cute::get<{i}>(kv_tile_shape))" + for i in range(self.na_dim) + ] + ) + dispatcher_str += ( + ' throw std::runtime_error("Blackwell FNA-' + + f"{self.na_dim}D forward kernel dispatch failed! " + f"It got invalid Q tile and KV tile combination " + f"({self.dtype.short_name}, head_dim {self.head_dim}): " + 'q_tile=(" + ' + q_tile_str + ' + "), ' + 'kv_tile=(" + ' + kv_tile_str + ' + ")."' + "); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +def write_header_file(content, path, namespaces, extra_includes=None): + extra_includes = extra_includes or [] + header_head = [ + "#pragma once\n", + "\n\n", + ] + header_head += ["#include \n"] + header_head += ["#include \n"] + header_head += ["#include \n"] + header_head += ["#include \n"] + header_head += ["#ifdef NATTEN_WITH_CUTLASS\n"] + header_head += ["#ifdef NATTEN_WITH_BLACKWELL_FNA\n"] + for incl in extra_includes: + header_head += [f"#include <{incl}> \n"] + + for namespace in namespaces: + header_head += [f"namespace {namespace}", " { \n"] + + header_foot = [ + "\n\n", + ] + for namespace in namespaces: + header_foot += ["} ", f"// namespace {namespace}", " \n"] + header_foot += [ + "#endif \n", + "#endif \n", + "\n", + ] + with open(path, "w") as f: + f.write("".join(header_head)) + f.write(content) + f.write("".join(header_foot)) + + +def generate_blackwell_fna_kernels(path, num_splits=2): + + NA_DIMS = [1, 2, 3] + + SUPPORTED_DTYPES = [ + Half, + BFloat, + E4M3, + E5M2, + ] + + HEAD_DIMS = [32, 64, 128] + + CAUSAL_MASKS = { + 1: [(False,), (True,)], + 2: [(False, False), (False, True), (True, False), (True, True)], + 3: [ + (False, False, False), + (False, False, True), + (False, True, False), + (False, True, True), + (True, False, False), + (True, False, True), + (True, True, False), + (True, True, True), + ], + } + + TILE_SHAPES = { + 1: [ + ((256,), (128,)), + ], + 2: [ + ((16, 16), (16, 8)), + ((16, 16), (8, 16)), + ((8, 32), (8, 16)), + ((8, 32), (4, 32)), + ], + 3: [ + ((8, 4, 8), (4, 4, 8)), + ((8, 4, 8), (2, 8, 8)), + ((2, 8, 16), (4, 4, 8)), + ((2, 8, 16), (2, 8, 8)), + ((4, 4, 16), (2, 4, 16)), + ((2, 16, 8), (2, 8, 8)), + ((4, 8, 8), (2, 8, 8)), + ], + } + + dtype_dispatchers = [] + head_dim_dispatchers = [] + cm_dispatchers = [] + tile_shape_dispatchers = [] + kernels = [] + + rank_dispatcher = NaDimDispatcher() + for na_dim in NA_DIMS: + rank_dispatcher.append(na_dim) + + dtype_dispatcher = DTypeDispatcher(na_dim=na_dim) + for dtype in SUPPORTED_DTYPES: + dtype_dispatcher.append(dtype) + + head_dim_dispatcher = HeadDimDispatcher(na_dim=na_dim, dtype=dtype) + + for head_dim in HEAD_DIMS: + head_dim_dispatcher.append(head_dim) + + cm_dispatcher = CausalMaskDispatcher( + na_dim=na_dim, dtype=dtype, head_dim=head_dim + ) + for cm in CAUSAL_MASKS[na_dim]: + cm_dispatcher.append(cm) + + tile_shape_dispatcher = TileShapeDispatcher( + na_dim=na_dim, + dtype=dtype, + head_dim=head_dim, + causal_mask=cm, + ) + + for tile_shape in TILE_SHAPES[na_dim]: + tile_shape_dispatcher.append(tile_shape) + for persistent in [False, True]: + kernels.append( + tile_shape_dispatcher.get_kernel_instance( + tile_shape[0], tile_shape[1], persistent + ) + ) + tile_shape_dispatchers.append(tile_shape_dispatcher) + cm_dispatchers.append(cm_dispatcher) + head_dim_dispatchers.append(head_dim_dispatcher) + dtype_dispatchers.append(dtype_dispatcher) + + # + + path_to_sources = f"{path}/autogen/src/cuda/blackwell_fna/" + rel_header = "natten_autogen/cuda/blackwell_fna/" + path_to_header_dir = f"{path}/autogen/include/{rel_header}" + + os.makedirs(path_to_sources, exist_ok=False) + os.makedirs(path_to_header_dir, exist_ok=False) + + path_headers = f"{path_to_header_dir}kernels.h" + path_rank = f"{path_to_header_dir}interface.h" + path_dtype = f"{path_to_header_dir}dispatch_dtype.h" + path_head_dim = f"{path_to_header_dir}dispatch_head_dim.h" + path_cm = f"{path_to_header_dir}dispatch_cm.h" + path_tile_shape = f"{path_to_header_dir}dispatch_tile_shape.h" + + rel_path_headers = f"{rel_header}kernels.h" + rel_path_dtype = f"{rel_header}dispatch_dtype.h" + rel_path_head_dim = f"{rel_header}dispatch_head_dim.h" + rel_path_cm = f"{rel_header}dispatch_cm.h" + rel_path_tile_shape = f"{rel_header}dispatch_tile_shape.h" + + rank_disp = rank_dispatcher.get_dispatcher() + + dtype_disp = "" + for dispatcher in dtype_dispatchers: + dtype_disp += dispatcher.get_dispatcher() + + head_dim_disp = "" + for dispatcher in head_dim_dispatchers: + head_dim_disp += dispatcher.get_dispatcher() + + cm_disp = "" + for dispatcher in cm_dispatchers: + cm_disp += dispatcher.get_dispatcher() + + tile_shape_disp = "" + for dispatcher in tile_shape_dispatchers: + tile_shape_disp += dispatcher.get_dispatcher() + + headers = "" + for kernel in kernels: + headers += kernel.get_decl() + + assert ( + len(kernels) >= num_splits + ), f"Generated {len(kernels)} kernels, but got {num_splits=}." + # print(f"Generated {len(kernels)} kernels, got {num_splits=}.") + split_size = len(kernels) // num_splits + num_splits_with_res = len(kernels) % num_splits + kernels_emitted = [] + kernels_split = [] + for split_idx in range(num_splits): + kernel_start_idx = split_size * split_idx + min(num_splits_with_res, split_idx) + num_kernels_in_split = split_size + ( + 1 if split_idx < num_splits_with_res else 0 + ) + kernel_end_idx = kernel_start_idx + num_kernels_in_split + assert kernel_end_idx <= len(kernels) + pth_set = set() + source_list = [] + for kernel_idx in range(kernel_start_idx, kernel_end_idx): + kernel = kernels[kernel_idx] + # pth_set.add(kernel.path_to_header) + source_list.append(kernel) + kernels_emitted.append(kernel_idx) + pth_set.add(rel_path_headers) + write_combined_source_file( + path_to_sources, f"source_{split_idx}.cu", sorted(pth_set), source_list + ) + kernels_split.append(source_list) + # print(f"{split_idx=}, {kernel_start_idx=}, {kernel_end_idx=}, {len(kernels_emitted)=}") + assert split_idx == num_splits - 1, f"Expected {split_idx=} == {num_splits=} - 1" + assert len(kernels_emitted) == len(kernels) + assert sorted(kernels_emitted) == [ + x for x in range(len(kernels)) + ], f"{sorted(kernels_emitted)=}" + assert all(len(x) > 0 for x in kernels_split) + + namespaces = ["natten", "cuda", "fna_blackwell"] + cuda_headers = [ + "natten/natten.h", + "ATen/ATen.h", + "ATen/cuda/CUDAContext.h", + "c10/cuda/CUDAGuard.h", + "c10/cuda/CUDAStream.h", + "natten/natten.h", + "natten/helpers.h", + "natten/cuda/fna_blackwell/fna_forward.cuh", + ] + write_header_file(rank_disp, path_rank, namespaces, cuda_headers + [rel_path_dtype]) + write_header_file( + dtype_disp, path_dtype, namespaces, cuda_headers + [rel_path_head_dim] + ) + write_header_file( + head_dim_disp, path_head_dim, namespaces, cuda_headers + [rel_path_cm] + ) + write_header_file( + cm_disp, path_cm, namespaces, cuda_headers + [rel_path_tile_shape] + ) + write_header_file( + tile_shape_disp, path_tile_shape, namespaces, cuda_headers + [rel_path_headers] + ) + write_header_file(headers, path_headers, namespaces, cuda_headers) + + +def generate_blackwell_fna(output_directory: str, num_splits: int): + generate_blackwell_fna_kernels(output_directory, num_splits=num_splits) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "-o", + "--output-directory", + default=DEFAULT_OUTPUT_DIR, + help="Path to the directory where the auto-generated " + "kernel instantiations are dumped. " + f"Default: {DEFAULT_OUTPUT_DIR}", + ) + parser.add_argument( + "--num-splits", + type=int, + default=56, + help="Number of source files into which the kernels are split. Default: 56.", + ) + args = parser.parse_args() + generate_blackwell_fna(args.output_directory, args.num_splits) diff --git a/natten/scripts/autogen_blackwell_fna_bwd.py b/natten/scripts/autogen_blackwell_fna_bwd.py new file mode 100644 index 00000000..08cbfbbd --- /dev/null +++ b/natten/scripts/autogen_blackwell_fna_bwd.py @@ -0,0 +1,790 @@ +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# This script is intended to emit fused kernel instantiations into +# a variable number of source files, generate appropriate headers +# and a single dispatcher interface, which will be used by the +# NATTEN API to call the kernels. +# +# NOTE: these scripts are heavily under-documented, and +# overly-repetitive, and will be replaced in future PRs. +# Please use it with caution. + +import argparse +import math +import os +from typing import List, Tuple + +DEFAULT_OUTPUT_DIR = "csrc/" + +SUPPORTED_GEMM_SHAPES_BACKWARD = [ + (128, 128, 32), + (128, 128, 64), + (128, 128, 128), +] + + +KERNEL_DECL_TEMPLATE = """ +void {kernel_name}( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + {DimType} q_shape, + {DimType} kv_shape, + {DimType} qkv_shape, + {DimType} window_size, + {DimType} stride, + {DimType} dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); +""" + + +KERNEL_IMPL_TEMPLATE = """ +void {kernel_name}( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads_q, + int heads_kv, + int dim, + float attn_scale, + // fna parameters + {DimType} q_shape, + {DimType} kv_shape, + {DimType} qkv_shape, + {DimType} window_size, + {DimType} stride, + {DimType} dilation, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) {{ + + using Causal = {Causal}; + using QTileShape = {QTileShape}; + using KVTileShape = {KVTileShape}; + using Config = {GEMMShape}; + using Kernel = natten::cuda::fna_blackwell::KernelBackward< + {dtype}, Causal, QTileShape, KVTileShape, Config>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads_q, + heads_kv, + dim, + attn_scale, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({{bytes}}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +}} +""" + + +class DataType: + def __init__(self, name, short_name, torch_name, bits): + self.name = name + self.bits = bits + self.short_name = short_name + self.torch_name = torch_name + + +Half = DataType("cutlass::half_t", "float16", "torch::kFloat16", 16) +BFloat = DataType("cutlass::bfloat16_t", "bfloat16", "torch::kBFloat16", 16) + + +def iterable_to_static_cute_tuple(shape_in) -> str: + shape = ", ".join([f"cute::Int<{x}>" for x in shape_in]) + return f"cute::tuple<{shape}>" + + +def get_dim_type(na_dim: int) -> str: + shape = ", ".join(["int" for _ in range(na_dim)]) + return f"cute::tuple<{shape}>" + + +class BlackwellFnaInstance: + def __init__( + self, + na_dim: int, + dtype: DataType, + gemm_shape: Tuple[int, int, int], + q_tile_shape: tuple, + kv_tile_shape: tuple, + causal: tuple, + ): + assert 0 < na_dim <= 3 + assert len(gemm_shape) == 3 + assert na_dim == len(q_tile_shape) + assert na_dim == len(kv_tile_shape) + assert na_dim == len(causal) + self.na_dim = na_dim + self.gemm_shape = gemm_shape + self.q_tile_shape = q_tile_shape + self.kv_tile_shape = kv_tile_shape + self.causal = causal + self.dtype = dtype + self.max_head_dim = gemm_shape[2] + + def get_q_tile_shape_cute(self) -> str: + return iterable_to_static_cute_tuple(self.q_tile_shape) + + def get_kv_tile_shape_cute(self) -> str: + return iterable_to_static_cute_tuple(self.kv_tile_shape) + + def get_gemm_shape_cute(self) -> str: + return iterable_to_static_cute_tuple(self.gemm_shape) + + def get_causal_cute(self) -> str: + consts = ", ".join( + ["cute::true_type" if c else "cute::false_type" for c in self.causal] + ) + return f"cute::tuple<{consts}>" + + def get_name(self) -> str: + name = f"blackwell_fna{self.na_dim}d_backward" + name += f"_{self.dtype.short_name}" + name += "_" + "x".join([str(x) for x in self.gemm_shape]) + name += "_Q" + "x".join([str(x) for x in self.q_tile_shape]) + name += "_KV" + "x".join([str(x) for x in self.kv_tile_shape]) + name += "_causal" + "x".join(["1" if c else "0" for c in self.causal]) + return name + + def get_decl(self) -> str: + return KERNEL_DECL_TEMPLATE.format( + kernel_name=self.get_name(), DimType=get_dim_type(self.na_dim) + ) + + def get_impl(self) -> str: + return KERNEL_IMPL_TEMPLATE.format( + kernel_name=self.get_name(), + DimType=get_dim_type(self.na_dim), + Causal=self.get_causal_cute(), + QTileShape=self.get_q_tile_shape_cute(), + KVTileShape=self.get_kv_tile_shape_cute(), + GEMMShape=self.get_gemm_shape_cute(), + dtype=self.dtype.name, + ) + + +def write_combined_source_file(path, filename, headers, kernels): + source_head = [] + source_head += ["#ifdef NATTEN_WITH_CUTLASS\n"] + source_head += ["#ifdef NATTEN_WITH_BLACKWELL_FNA\n"] + + source_head += ["#include \n"] + source_head += ["#include \n"] + + source_head += ["#include \n"] + source_head += ["#include \n"] + source_head += ["#include \n"] + source_head += ["#include \n"] + + source_head += ["#include \n"] + source_head += ["#include \n"] + + source_head += ["#include \n"] + + for header in headers: + source_head += [f"#include <{header}>\n"] + + source_head += ["namespace natten { \n"] + source_head += ["namespace cuda { \n"] + source_head += ["namespace fna_blackwell { \n\n"] + + source_head = "".join(source_head) + + source_body = [] + for kernel in kernels: + source_body += "\n\n" + kernel.get_impl() + "\n\n" + source_body = "".join(source_body) + + source_foot = "".join( + [ + "} // namespace fna_blackwell \n", + "} // namespace cuda \n", + "} // namespace natten \n", + "#endif \n", + "#endif \n", + "\n", + ] + ) + filename = f"{path}/{filename}" + with open(filename, "w") as f: + f.write(source_head) + f.write(source_body) + f.write(source_foot) + + +class NaDimDispatcher: + def __init__(self): + self.name = "DISPATCH_BLACKWELL_FNA_BACKWARD" + self.dims = [] + + def append(self, na_dim: int): + self.dims.append(na_dim) + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += f"#define {self.name}(rank, dtype, dim, is_causal, q_tile_shape, kv_tile_shape, ...) \\\n" + dispatcher_str += " [&] { \\\n" + for i, na_dim in enumerate(self.dims): + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + dispatcher_str += f"if constexpr (rank == {na_dim})" + dispatcher_str += " { \\\n" + dispatcher_str += " " + dispatcher_str += f" {self.name}_{na_dim}D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \\\n" + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += ( + ' throw std::runtime_error("Blackwell FNA backward kernel dispatch failed! ' + + "It only supports NA1D, 2D, and 3D!" + + '"); \\\n' + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class DTypeDispatcher: + def __init__(self, na_dim: int): + self.dtypes: List[DataType] = [] + self.na_dim = na_dim + self.name = f"DISPATCH_BLACKWELL_FNA_BACKWARD_{self.na_dim}D" + + def append(self, dtype: DataType): + self.dtypes.append(dtype) + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += f"#define {self.name}(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, ...) \\\n" + dispatcher_str += " [&] { \\\n" + for i, dtype in enumerate(self.dtypes): + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + dispatcher_str += f"if (dtype == {dtype.torch_name})" + dispatcher_str += " { \\\n" + dispatcher_str += " " + dispatcher_str += f" {self.name}_{dtype.short_name}(dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \\\n" + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += ( + ' throw std::runtime_error("Blackwell FNA-' + + f"{self.na_dim}D backward kernel dispatch failed! " + 'It does not support dtype "' + ' + std::string(c10::toString(dtype)) + "."' + "); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class HeadDimDispatcher: + def __init__(self, na_dim: int, dtype: DataType): + self.na_dim = na_dim + self.dtype = dtype + self.name = ( + f"DISPATCH_BLACKWELL_FNA_BACKWARD_{self.na_dim}D_{self.dtype.short_name}" + ) + + self.dims: List[int] = [] + + def append(self, dim: int): + self.dims.append(dim) + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += f"#define {self.name}(dim, is_causal, q_tile_shape, kv_tile_shape, ...) \\\n" + dispatcher_str += " [&] { \\\n" + for i, dim in enumerate(sorted(self.dims)): + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + dispatcher_str += f"if (dim <= {dim})" + dispatcher_str += " { \\\n" + dispatcher_str += " " + dispatcher_str += f" {self.name}_headdim{dim}(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \\\n" + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += ( + " throw std::runtime_error(" + '"Blackwell FNA backward kernel dispatch failed! ' + 'It does not support head dim "' + + f' + std::to_string(dim) + " for {self.dtype.short_name}."' + + "); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class CausalMaskDispatcher: + def __init__( + self, + na_dim: int, + dtype: DataType, + head_dim: int, + ): + self.na_dim = na_dim + self.dtype = dtype + self.head_dim = head_dim + + self.name = f"DISPATCH_BLACKWELL_FNA_BACKWARD_{self.na_dim}D_{self.dtype.short_name}_headdim{head_dim}" + + self.cms: List = [] + + def append(self, cm): + assert len(cm) == self.na_dim, f"{cm} incompatible for {self.na_dim}D NA." + self.cms.append(cm) + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += ( + f"#define {self.name}(is_causal, q_tile_shape, kv_tile_shape, ...) \\\n" + ) + dispatcher_str += " [&] { \\\n" + i = 0 + for cm in self.cms: + cm_str = "causal" + "x".join(["1" if c else "0" for c in cm]) + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + i += 1 + dispatcher_str += "if (" + for dim in range(self.na_dim): + dispatcher_str += ( + f"cute::get<{dim}>(is_causal)" + if cm[dim] + else f"not cute::get<{dim}>(is_causal)" + ) + if dim != self.na_dim - 1: + dispatcher_str += " && " + dispatcher_str += ")" + dispatcher_str += " { \\\n" + dispatcher_str += " " + dispatcher_str += f" {self.name}_{cm_str}(q_tile_shape, kv_tile_shape, __VA_ARGS__); \\\n" + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += " " + dispatcher_str += ( + ' throw std::runtime_error("Blackwell FNA-' + + f"{self.na_dim}D backward kernel dispatch failed! " + f"Causal mask dispatcher ({self.dtype.short_name}, " + f'head_dim {self.head_dim}) got invalid causal mask!"' + "); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class TileShapeDispatcher: + def __init__( + self, + na_dim: int, + dtype: DataType, + head_dim: int, + causal_mask: tuple, + ): + self.na_dim = na_dim + self.dtype = dtype + self.head_dim = head_dim + self.causal_mask = causal_mask + cm_str = "_causal" + "x".join(["1" if c else "0" for c in self.causal_mask]) + + self.name = f"DISPATCH_BLACKWELL_FNA_BACKWARD_{self.na_dim}D_{self.dtype.short_name}_headdim{head_dim}" + self.name += cm_str + + self.tile_shapes: List = [] + + def append(self, tile_shape): + assert len(tile_shape) == 2 + assert len(tile_shape[0]) == self.na_dim + assert len(tile_shape[1]) == self.na_dim + self.tile_shapes.append(tile_shape) + + def get_kernel_instance(self, q_tile_shape, kv_tile_shape): + gemm_M = math.prod(q_tile_shape) + gemm_N = math.prod(kv_tile_shape) + gemm_K = self.head_dim + gemm_shape = (gemm_M, gemm_N, gemm_K) + + assert gemm_shape in SUPPORTED_GEMM_SHAPES_BACKWARD + + kernel = BlackwellFnaInstance( + na_dim=self.na_dim, + dtype=self.dtype, + gemm_shape=gemm_shape, + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + causal=self.causal_mask, + ) + + return kernel + + def get_target_name(self, q_tile_shape, kv_tile_shape): + kernel = self.get_kernel_instance(q_tile_shape, kv_tile_shape) + return kernel.get_name() + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += f"#define {self.name}(q_tile_shape, kv_tile_shape, ...) \\\n" + dispatcher_str += " [&] { \\\n" + i = 0 + for q_tile_shape, kv_tile_shape in self.tile_shapes: + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + i += 1 + dispatcher_str += "if (" + for dim in range(self.na_dim): + dispatcher_str += ( + f"cute::get<{dim}>(q_tile_shape) == {q_tile_shape[dim]}" + ) + dispatcher_str += " && " + dispatcher_str += "\\\n" + for dim in range(self.na_dim): + dispatcher_str += ( + f"cute::get<{dim}>(kv_tile_shape) == {kv_tile_shape[dim]}" + ) + if dim != self.na_dim - 1: + dispatcher_str += " && " + dispatcher_str += ")" + dispatcher_str += " { \\\n" + + dispatcher_str += f" natten::cuda::fna_blackwell::{self.get_target_name(q_tile_shape, kv_tile_shape)}(__VA_ARGS__); \\\n" + + dispatcher_str += "} \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += " " + q_tile_str = ' + "," + '.join( + [ + f"std::to_string(cute::get<{i}>(q_tile_shape))" + for i in range(self.na_dim) + ] + ) + kv_tile_str = ' + "," + '.join( + [ + f"std::to_string(cute::get<{i}>(kv_tile_shape))" + for i in range(self.na_dim) + ] + ) + dispatcher_str += ( + ' throw std::runtime_error("Blackwell FNA-' + + f"{self.na_dim}D backward kernel dispatch failed! " + f"It got invalid Q tile and KV tile combination " + f"({self.dtype.short_name}, head_dim {self.head_dim}): " + 'q_tile=(" + ' + q_tile_str + ' + "), ' + 'kv_tile=(" + ' + kv_tile_str + ' + ")."' + "); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +def write_header_file(content, path, namespaces, extra_includes=None): + extra_includes = extra_includes or [] + header_head = [ + "#pragma once\n", + "\n\n", + ] + header_head += ["#include \n"] + header_head += ["#include \n"] + header_head += ["#include \n"] + header_head += ["#include \n"] + header_head += ["#ifdef NATTEN_WITH_CUTLASS\n"] + header_head += ["#ifdef NATTEN_WITH_BLACKWELL_FNA\n"] + for incl in extra_includes: + header_head += [f"#include <{incl}> \n"] + + for namespace in namespaces: + header_head += [f"namespace {namespace}", " { \n"] + + header_foot = [ + "\n\n", + ] + for namespace in namespaces: + header_foot += ["} ", f"// namespace {namespace}", " \n"] + header_foot += [ + "#endif \n", + "#endif \n", + "\n", + ] + with open(path, "w") as f: + f.write("".join(header_head)) + f.write(content) + f.write("".join(header_foot)) + + +def generate_blackwell_fna_kernels(path, num_splits=2): + + NA_DIMS = [1, 2, 3] + + SUPPORTED_DTYPES = [ + Half, + BFloat, + ] + + HEAD_DIMS = [32, 64, 128] + + CAUSAL_MASKS = { + 1: [(False,), (True,)], + 2: [(False, False), (False, True), (True, False), (True, True)], + 3: [ + (False, False, False), + (False, False, True), + (False, True, False), + (False, True, True), + (True, False, False), + (True, False, True), + (True, True, False), + (True, True, True), + ], + } + + TILE_SHAPES = { + 1: [ + ((128,), (128,)), + ], + 2: [ + ((16, 8), (16, 8)), + ((16, 8), (8, 16)), + ((8, 16), (16, 8)), + ((8, 16), (8, 16)), + ], + 3: [ + ((4, 4, 8), (4, 4, 8)), + ((4, 4, 8), (2, 8, 8)), + ((1, 8, 16), (4, 4, 8)), + ((2, 8, 8), (4, 4, 8)), + ((1, 8, 16), (2, 8, 8)), + ((2, 4, 16), (2, 4, 16)), + ((4, 2, 16), (2, 4, 16)), + ((4, 4, 8), (2, 4, 16)), + ((2, 8, 8), (2, 8, 8)), + ], + } + + dtype_dispatchers = [] + head_dim_dispatchers = [] + cm_dispatchers = [] + tile_shape_dispatchers = [] + kernels = [] + + rank_dispatcher = NaDimDispatcher() + for na_dim in NA_DIMS: + rank_dispatcher.append(na_dim) + + dtype_dispatcher = DTypeDispatcher(na_dim=na_dim) + for dtype in SUPPORTED_DTYPES: + dtype_dispatcher.append(dtype) + + head_dim_dispatcher = HeadDimDispatcher(na_dim=na_dim, dtype=dtype) + + for head_dim in HEAD_DIMS: + head_dim_dispatcher.append(head_dim) + + cm_dispatcher = CausalMaskDispatcher( + na_dim=na_dim, dtype=dtype, head_dim=head_dim + ) + for cm in CAUSAL_MASKS[na_dim]: + cm_dispatcher.append(cm) + + tile_shape_dispatcher = TileShapeDispatcher( + na_dim=na_dim, + dtype=dtype, + head_dim=head_dim, + causal_mask=cm, + ) + + for tile_shape in TILE_SHAPES[na_dim]: + tile_shape_dispatcher.append(tile_shape) + kernels.append( + tile_shape_dispatcher.get_kernel_instance( + tile_shape[0], tile_shape[1] + ) + ) + tile_shape_dispatchers.append(tile_shape_dispatcher) + cm_dispatchers.append(cm_dispatcher) + head_dim_dispatchers.append(head_dim_dispatcher) + dtype_dispatchers.append(dtype_dispatcher) + + # + + path_to_sources = f"{path}/autogen/src/cuda/blackwell_fna_bwd/" + rel_header = "natten_autogen/cuda/blackwell_fna_bwd/" + path_to_header_dir = f"{path}/autogen/include/{rel_header}" + + os.makedirs(path_to_sources, exist_ok=False) + os.makedirs(path_to_header_dir, exist_ok=False) + + path_headers = f"{path_to_header_dir}kernels.h" + path_rank = f"{path_to_header_dir}interface.h" + path_dtype = f"{path_to_header_dir}dispatch_dtype.h" + path_head_dim = f"{path_to_header_dir}dispatch_head_dim.h" + path_cm = f"{path_to_header_dir}dispatch_cm.h" + path_tile_shape = f"{path_to_header_dir}dispatch_tile_shape.h" + + rel_path_headers = f"{rel_header}kernels.h" + rel_path_dtype = f"{rel_header}dispatch_dtype.h" + rel_path_head_dim = f"{rel_header}dispatch_head_dim.h" + rel_path_cm = f"{rel_header}dispatch_cm.h" + rel_path_tile_shape = f"{rel_header}dispatch_tile_shape.h" + + rank_disp = rank_dispatcher.get_dispatcher() + + dtype_disp = "" + for dispatcher in dtype_dispatchers: + dtype_disp += dispatcher.get_dispatcher() + + head_dim_disp = "" + for dispatcher in head_dim_dispatchers: + head_dim_disp += dispatcher.get_dispatcher() + + cm_disp = "" + for dispatcher in cm_dispatchers: + cm_disp += dispatcher.get_dispatcher() + + tile_shape_disp = "" + for dispatcher in tile_shape_dispatchers: + tile_shape_disp += dispatcher.get_dispatcher() + + headers = "" + for kernel in kernels: + headers += kernel.get_decl() + + assert ( + len(kernels) >= num_splits + ), f"Generated {len(kernels)} kernels, but got {num_splits=}." + split_size = len(kernels) // num_splits + num_splits_with_res = len(kernels) % num_splits + kernels_emitted = [] + kernels_split = [] + for split_idx in range(num_splits): + kernel_start_idx = split_size * split_idx + min(num_splits_with_res, split_idx) + num_kernels_in_split = split_size + ( + 1 if split_idx < num_splits_with_res else 0 + ) + kernel_end_idx = kernel_start_idx + num_kernels_in_split + assert kernel_end_idx <= len(kernels) + pth_set = set() + source_list = [] + for kernel_idx in range(kernel_start_idx, kernel_end_idx): + kernel = kernels[kernel_idx] + # pth_set.add(kernel.path_to_header) + source_list.append(kernel) + kernels_emitted.append(kernel_idx) + pth_set.add(rel_path_headers) + write_combined_source_file( + path_to_sources, f"source_{split_idx}.cu", sorted(pth_set), source_list + ) + kernels_split.append(source_list) + # print(f"{split_idx=}, {kernel_start_idx=}, {kernel_end_idx=}, {len(kernels_emitted)=}") + assert split_idx == num_splits - 1, f"Expected {split_idx=} == {num_splits=} - 1" + assert len(kernels_emitted) == len(kernels) + assert sorted(kernels_emitted) == [ + x for x in range(len(kernels)) + ], f"{sorted(kernels_emitted)=}" + assert all(len(x) > 0 for x in kernels_split) + + namespaces = ["natten", "cuda", "fna_blackwell"] + cuda_headers = [ + "natten/natten.h", + "ATen/ATen.h", + "ATen/cuda/CUDAContext.h", + "c10/cuda/CUDAGuard.h", + "c10/cuda/CUDAStream.h", + "natten/natten.h", + "natten/helpers.h", + "natten/cuda/fna_blackwell/fna_backward.cuh", + ] + write_header_file(rank_disp, path_rank, namespaces, cuda_headers + [rel_path_dtype]) + write_header_file( + dtype_disp, path_dtype, namespaces, cuda_headers + [rel_path_head_dim] + ) + write_header_file( + head_dim_disp, path_head_dim, namespaces, cuda_headers + [rel_path_cm] + ) + write_header_file( + cm_disp, path_cm, namespaces, cuda_headers + [rel_path_tile_shape] + ) + write_header_file( + tile_shape_disp, path_tile_shape, namespaces, cuda_headers + [rel_path_headers] + ) + write_header_file(headers, path_headers, namespaces, cuda_headers) + + +def generate_blackwell_fna(output_directory: str, num_splits: int): + generate_blackwell_fna_kernels(output_directory, num_splits=num_splits) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "-o", + "--output-directory", + default=DEFAULT_OUTPUT_DIR, + help="Path to the directory where the auto-generated " + "kernel instantiations are dumped. " + f"Default: {DEFAULT_OUTPUT_DIR}", + ) + parser.add_argument( + "--num-splits", + type=int, + default=28, + help="Number of source files into which the kernels are split. Default: 28.", + ) + args = parser.parse_args() + generate_blackwell_fna(args.output_directory, args.num_splits) diff --git a/natten/scripts/autogen_fmha.py b/natten/scripts/autogen_fmha.py new file mode 100644 index 00000000..d30b4062 --- /dev/null +++ b/natten/scripts/autogen_fmha.py @@ -0,0 +1,676 @@ +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# This script is intended to emit fused kernel instantiations into +# a variable number of source files, generate appropriate headers +# and a single dispatcher interface, which will be used by the +# NATTEN API to call the kernels. +# +# NOTE: these scripts are heavily under-documented, and +# overly-repetitive, and will be replaced in future PRs. +# Please use it with caution. + +import argparse +import os +from typing import List, Tuple + +DEFAULT_OUTPUT_DIR = "csrc/" + +_SM_TO_UPPER_BOUND = { + 50: 70, + 70: 75, + 75: 80, + 80: None, +} + +_SM_TO_LOWER_BOUND = { + 50: 50, + 70: 70, + 75: 75, + 80: 80, +} + + +class KernelConfig: + def __init__(self, sm: int, gemm_shape: Tuple[int, int, int]): + assert 50 <= sm <= 80 + assert len(gemm_shape) == 3 + self.sm = sm + self.gemm_shape = gemm_shape + + def get_name(self, is_backward: bool) -> str: + backward_str = "" if not is_backward else "_backward" + return f"fmha{backward_str}_{self.gemm_shape[0]}x{self.gemm_shape[1]}x{self.gemm_shape[2]}_sm{self.sm}" + + +class KernelConfigList: + def __init__(self, sm: int, gemm_shapes: List[Tuple[int, int, int]]): + assert 50 <= sm <= 80 + self.sm = sm + self.gemm_shapes = gemm_shapes + + @property + def configs(self) -> List[KernelConfig]: + return [ + KernelConfig(sm=self.sm, gemm_shape=gemm_shape) + for gemm_shape in self.gemm_shapes + ] + + def get_name(self, is_backward: bool) -> str: + backward_str = "" if not is_backward else "_backward" + return f"fmha{backward_str}_sm{self.sm}" + + +class DataType: + def __init__(self, name, torch_name, short_name, bits, min_sm): + self.name = name + self.torch_name = torch_name + self.bits = bits + self.short_name = short_name + self.min_sm = min_sm + + +NATTEN_Float = DataType("float", "torch::kFloat32", "float32", 32, 0) +NATTEN_Half = DataType("cutlass::half_t", "torch::kFloat16", "float16", 16, 50) +NATTEN_BFloat = DataType("cutlass::bfloat16_t", "torch::kBFloat16", "bfloat16", 16, 80) + + +KERNEL_DECL_TEMPLATE = """__global__ void __launch_bounds__( + {CPP_CLASS}::kNumThreads, + {CPP_CLASS}::kMinBlocksPerSm) +{NAME}(typename {CPP_CLASS}::Params p); +""" + + +KERNEL_IMPL_TEMPLATE = """__global__ void __launch_bounds__( + {CPP_CLASS}::kNumThreads, + {CPP_CLASS}::kMinBlocksPerSm) +{NAME}(typename {CPP_CLASS}::Params p) {{ +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= {SM}0 +#if __CUDA_ARCH__ < {SM_MAX}0 + if (!p.advance_to_block()) {{ + return; + }} + {CPP_CLASS}::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: fmha kernel `{NAME}` was built for SM{SM}, but attempted to launch from SM%d\\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +}} +""" + +SM_TO_SM_MAX = { + 50: 70, + 70: 75, + 75: 80, + 80: 130, +} + + +class FusedNAKernel: + def __init__( + self, + is_backward: bool, + dtype: DataType, + config: KernelConfig, + ): + self.dtype = dtype + self.config = config + self.name_cc = f"{self.config.get_name(is_backward)}_{dtype.short_name}" + self.path_to_header = ( + "natten/cuda/fmha/kernel_forward.h" + if not is_backward + else "natten/cuda/fmha/kernel_backward.h" + ) + self.aligned = True + self.is_backward = is_backward + + @property + def cpp_class(self) -> str: + if self.is_backward: + template_args = ", ".join( + [ + self.dtype.name, + f"cutlass::arch::Sm{self.config.sm}", + "true" if self.aligned else "false", + # "false", # preload + str(self.config.gemm_shape[0]), + str(self.config.gemm_shape[1]), + str(self.config.gemm_shape[2]), + "false", # QKV divides tile sizes + ] + ) + return f"AttentionBackwardKernel<{template_args}>" + else: + template_args = ", ".join( + [ + self.dtype.name, + f"cutlass::arch::Sm{self.config.sm}", + "true" if self.aligned else "false", + str(self.config.gemm_shape[0]), + str(self.config.gemm_shape[1]), + str(self.config.gemm_shape[2]), + ] + ) + return f"AttentionKernel<{template_args}>" + + def header(self): + return KERNEL_DECL_TEMPLATE.format( + CPP_CLASS=self.cpp_class, + NAME=self.name_cc, + ) + + def source(self): + return KERNEL_IMPL_TEMPLATE.format( + CPP_CLASS=self.cpp_class, + NAME=self.name_cc, + SM=self.config.sm, + SM_MAX=SM_TO_SM_MAX[self.config.sm], + ) + + +class FusedNAKernelBundle: + def __init__( + self, + is_backward: bool, + sm: int, + dtype: DataType, + config_list: KernelConfigList, + ): + self.sm = sm + self.dtype = dtype + self.config_list = config_list + self.name_cc = f"{self.config_list.get_name(is_backward)}_{dtype.short_name}" + self.path_to_header = ( + "natten/cuda/fmha/kernel_forward.h" + if not is_backward + else "natten/cuda/fmha/kernel_backward.h" + ) + self.aligned = True + self.is_backward = is_backward + + @property + def kernels(self): + return [ + FusedNAKernel( + is_backward=self.is_backward, + dtype=self.dtype, + config=config, + ) + for config in self.config_list.configs + ] + + def header(self): + decl = "\n\n" + decl += "///////////////////////////////////////////////////////////////////\n" + decl += f"// FMHA / {self.dtype.short_name} / SM{self.sm}" + decl += "Backward Kernel" if self.is_backward else "" + decl += "\n" + decl += "///////////////////////////////////////////////////////////////////" + + for kernel in self.kernels: + decl += "\n\n" + decl += kernel.header() + + decl += f"\n\ntemplate \nvoid {self.name_cc}" + decl += "(T cb) {\n" + for kernel in self.kernels: + decl += f" cb({kernel.cpp_class}(), {kernel.name_cc});\n" + decl += "}" + return decl + + def source(self): + impl = "\n\n" + impl += "///////////////////////////////////////////////////////////////////\n" + impl += f"// FMHA / {self.dtype.short_name} / SM{self.sm}\n" + impl += "///////////////////////////////////////////////////////////////////" + + for kernel in self.kernels: + impl += "\n\n" + impl += kernel.source() + return impl + + +def write_combined_source_file(path, filename, headers, sources): + source_head = [] + source_head += ["#include \n"] + source_head += ["#include \n"] + + for header in headers: + source_head += [f"#include <{header}>\n"] + + source_head += ["namespace natten { \n"] + source_head += ["namespace cuda { \n"] + source_head += ["namespace fmha { \n\n"] + + source_head = "".join(source_head) + + source_body = [] + for source in sources: + source_body += source.source() + source_body = "".join(source_body) + + source_foot = "".join( + [ + "} \n", + "} \n", + "} \n", + "\n", + ] + ) + filename = f"{path}/{filename}" + with open(filename, "w") as f: + f.write(source_head) + f.write(source_body) + f.write(source_foot) + + +class DeviceDispatcher: + def __init__(self, is_backward: bool): + self.name_cc = ( + "DISPATCH_FMHA_FORWARD" if not is_backward else "DISPATCH_FMHA_BACKWARD" + ) + self.name_target = self.name_cc + "_SM" + self.devices: List[int] = [] + self.is_backward = is_backward + + def append(self, sm: int): + self.devices.append(sm) + + def get_dispatcher(self): + dispatcher_str = "" + if self.is_backward: + dispatcher_str += f"#define {self.name_cc}(cc, dtype, cb) \\\n" + else: + dispatcher_str += f"#define {self.name_cc}(cc, dtype, cb) \\\n" + dispatcher_str += " [&] { \\\n" + for i, sm in enumerate(self.devices): + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + dispatcher_str += "if (" + if _SM_TO_UPPER_BOUND[sm] is not None: + dispatcher_str += f"cc < {_SM_TO_UPPER_BOUND[sm]} && " + dispatcher_str += f"cc >= {_SM_TO_LOWER_BOUND[sm]})" + dispatcher_str += " { \\\n" + dispatcher_str += " " + if self.is_backward: + dispatcher_str += f" {self.name_target}{sm}(dtype, cb); \\\n" + else: + dispatcher_str += f" {self.name_target}{sm}(dtype, cb); \\\n" + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += ( + ' throw std::runtime_error("NATTEN FMHA kernel dispatch failed! ' + + "FMHA is not implemented for this device." + + '"); \\\n' + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class DataTypeDispatcher: + def __init__(self, is_backward: bool, sm: int): + self.dtypes: List[DataType] = [] + self.sm = sm + self.name_cc = ( + f"DISPATCH_FMHA_FORWARD_SM{self.sm}" + if not is_backward + else f"DISPATCH_FMHA_BACKWARD_SM{self.sm}" + ) + backward_str = "" if not is_backward else "_backward" + self.name_target = f"fmha{backward_str}_sm{self.sm}" + self.is_backward = is_backward + + def append(self, dtype: DataType): + self.dtypes.append(dtype) + + def get_dispatcher(self): + dispatcher_str = "" + if self.is_backward: + dispatcher_str += f"#define {self.name_cc}(dtype, cb) \\\n" + else: + dispatcher_str += f"#define {self.name_cc}(dtype, cb) \\\n" + dispatcher_str += " [&] { \\\n" + for i, dtype in enumerate(self.dtypes): + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + dispatcher_str += f"if (dtype == {dtype.torch_name})" + dispatcher_str += " { \\\n" + dispatcher_str += " " + if self.is_backward: + dispatcher_str += f" {self.name_target}_{dtype.short_name}(cb); \\\n" + else: + dispatcher_str += f" {self.name_target}_{dtype.short_name}(cb); \\\n" + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += ( + ' throw std::runtime_error("NATTEN kernel dispatch failed! ' + + f"FMHA does not support this data type on SM{self.sm}." + + '"); \\\n' + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +def write_header_file(content, path, namespaces, extra_includes=None): + extra_includes = extra_includes or [] + header_head = [ + "#pragma once\n", + "\n\n", + ] + header_head += ["#include \n"] + header_head += ["#include \n"] + header_head += ["#include \n"] + for incl in extra_includes: + header_head += [f"#include <{incl}> \n"] + + for namespace in namespaces: + header_head += [f"namespace {namespace}", " { \n"] + + header_foot = [ + "\n\n", + ] + for namespace in namespaces: + header_foot += ["} ", f"// namespace {namespace}", " \n"] + header_foot += [ + "\n", + ] + with open(path, "w") as f: + f.write("".join(header_head)) + f.write(content) + f.write("".join(header_foot)) + + +def generate_cuda_kernels(path, num_splits=2): + + SUPPORTED_ARCHS = [50, 70, 75, 80] + + CUDA_DTYPES = [ + NATTEN_Float, + NATTEN_Half, + NATTEN_BFloat, + ] + + GEMM_SHAPES = [ + # K = 32 + (32, 128, 32), + (64, 64, 32), + (64, 128, 32), + # K = 64 + (32, 128, 64), + (64, 64, 64), + (64, 128, 64), + # K = 128 + (32, 128, 128), + (64, 64, 128), + (64, 128, 128), + # K = 65536 + (32, 128, 2**16), + (64, 128, 2**16), + (64, 64, 2**16), + ] + dtype_dispatchers = [] + kernels = [] + + device_dispatcher = DeviceDispatcher(is_backward=False) + for sm in SUPPORTED_ARCHS: + device_dispatcher.append(sm) + dtype_dispatcher = DataTypeDispatcher(is_backward=False, sm=sm) + for dtype in CUDA_DTYPES: + if dtype.min_sm > sm: + continue + dtype_dispatcher.append(dtype) + kernels.append( + FusedNAKernelBundle( + is_backward=False, + sm=sm, + dtype=dtype, + config_list=KernelConfigList( + sm=sm, + gemm_shapes=GEMM_SHAPES, + ), + ) + ) + dtype_dispatchers.append(dtype_dispatcher) + + # + + BACKWARD_GEMM_SHAPES = { + 80: { + NATTEN_Half: [ + # K = 32 + (64, 64, 32), + # K = 64 + (64, 64, 64), + # K = 128 + (64, 64, 128), + (128, 128, 128), # SM80/SM90 only + # K = 65536 + (64, 64, 2**16), + (128, 64, 2**16), # SM80/SM90 only + ], + NATTEN_BFloat: [ + # K = 32 + (64, 64, 32), + # K = 64 + (64, 64, 64), + # K = 128 + (64, 64, 128), + (128, 128, 128), # SM80/SM90 only + # K = 65536 + (64, 64, 2**16), + (128, 64, 2**16), # SM80/SM90 only + ], + NATTEN_Float: [ + # K = 32 + (64, 64, 32), + # K = 64 + (64, 64, 64), + # K = 128 + (64, 64, 128), + (128, 64, 128), # SM80/SM90 only + # K = 65536 + (64, 64, 2**16), + (128, 64, 2**16), # SM80/SM90 only + ], + }, + 75: { + NATTEN_Half: [ + # K = 32 + (64, 64, 32), + # K = 64 + (64, 64, 64), + # K = 128 + (64, 64, 128), + # K = 65536 + (64, 64, 2**16), + ], + NATTEN_Float: [ + # K = 32 + (64, 64, 32), + # K = 64 + (64, 64, 64), + # K = 128 + (64, 64, 128), + # K = 65536 + (64, 64, 2**16), + ], + }, + 70: { + NATTEN_Half: [ + # K = 32 + (64, 64, 32), + # K = 64 + (64, 64, 64), + # K = 128 + (64, 64, 128), + (128, 64, 128), # SM70 only? + # K = 65536 + (64, 64, 2**16), + (128, 64, 2**16), # SM70 only? + ], + NATTEN_Float: [ + # K = 32 + (64, 64, 32), + # K = 64 + (64, 64, 64), + # K = 128 + (64, 64, 128), + # K = 65536 + (64, 64, 2**16), + ], + }, + 50: { + NATTEN_Half: [ + # K = 32 + (64, 64, 32), + # K = 64 + (64, 64, 64), + # K = 128 + (64, 64, 128), + # K = 65536 + (64, 64, 2**16), + ], + NATTEN_Float: [ + # K = 32 + (64, 64, 32), + # K = 64 + (64, 64, 64), + # K = 128 + (64, 64, 128), + # K = 65536 + (64, 64, 2**16), + ], + }, + } + + # Backward kernels + device_dispatcher_bwd = DeviceDispatcher(is_backward=True) + for sm in SUPPORTED_ARCHS: + device_dispatcher_bwd.append(sm) + dtype_dispatcher = DataTypeDispatcher(is_backward=True, sm=sm) + for dtype in CUDA_DTYPES: + if dtype.min_sm > sm: + continue + dtype_dispatcher.append(dtype) + kernels.append( + FusedNAKernelBundle( + is_backward=True, + sm=sm, + dtype=dtype, + config_list=KernelConfigList( + sm=sm, + gemm_shapes=BACKWARD_GEMM_SHAPES[sm][dtype], + ), + ) + ) + dtype_dispatchers.append(dtype_dispatcher) + + # + + path_to_sources = f"{path}/autogen/src/cuda/fmha/" + rel_header = "natten_autogen/cuda/fmha/" + path_to_header_dir = f"{path}/autogen/include/{rel_header}" + + os.makedirs(path_to_sources, exist_ok=False) + os.makedirs(path_to_header_dir, exist_ok=False) + + path_headers = f"{path_to_header_dir}kernels.h" + path_device = f"{path_to_header_dir}interface.h" + path_dtype = f"{path_to_header_dir}dispatch_dtype.h" + + rel_path_headers = f"{rel_header}kernels.h" + rel_path_dtype = f"{rel_header}dispatch_dtype.h" + + device_disp = ( + device_dispatcher.get_dispatcher() + device_dispatcher_bwd.get_dispatcher() + ) + + dtype_disp = "" + for dispatcher in dtype_dispatchers: + dtype_disp += dispatcher.get_dispatcher() + + headers = "" + for kernel in kernels: + headers += kernel.header() + + assert ( + len(kernels) >= num_splits + ), f"Generated {len(kernels)} kernels, but got {num_splits=}." + split_size = len(kernels) // num_splits + num_splits_with_res = len(kernels) % num_splits + kernels_emitted = [] + kernels_split = [] + for split_idx in range(num_splits): + kernel_start_idx = split_size * split_idx + min(num_splits_with_res, split_idx) + num_kernels_in_split = split_size + ( + 1 if split_idx < num_splits_with_res else 0 + ) + kernel_end_idx = kernel_start_idx + num_kernels_in_split + assert kernel_end_idx <= len(kernels) + pth_set = set() + source_list = [] + for kernel_idx in range(kernel_start_idx, kernel_end_idx): + kernel = kernels[kernel_idx] + pth_set.add(kernel.path_to_header) + source_list.append(kernel) + kernels_emitted.append(kernel_idx) + pth_set.add(rel_path_headers) + write_combined_source_file( + path_to_sources, f"source_{split_idx}.cu", sorted(pth_set), source_list + ) + kernels_split.append(source_list) + # print(f"{split_idx=}, {kernel_start_idx=}, {kernel_end_idx=}, {len(kernels_emitted)=}") + assert split_idx == num_splits - 1, f"Expected {split_idx=} == {num_splits=} - 1" + assert len(kernels_emitted) == len(kernels) + assert sorted(kernels_emitted) == [ + x for x in range(len(kernels)) + ], f"{sorted(kernels_emitted)=}" + assert all(len(x) > 0 for x in kernels_split) + + namespaces = ["natten", "cuda", "fmha"] + cuda_headers = [ + "natten/natten.h", + "natten/cuda/fmha/kernel_forward.h", + "natten/cuda/fmha/kernel_backward.h", + ] + write_header_file( + device_disp, path_device, namespaces, cuda_headers + [rel_path_dtype] + ) + write_header_file( + dtype_disp, path_dtype, namespaces, cuda_headers + [rel_path_headers] + ) + write_header_file(headers, path_headers, namespaces, cuda_headers) + + +def generate_cuda_fused(output_directory: str, num_splits: int): + generate_cuda_kernels(output_directory, num_splits=num_splits) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "-o", + "--output-directory", + default=DEFAULT_OUTPUT_DIR, + help="Path to the directory where the auto-generated " + "kernel instantiations are dumped. " + f"Default: {DEFAULT_OUTPUT_DIR}", + ) + parser.add_argument( + "--num-splits", + type=int, + default=12, + help="Number of source files into which the kernels are split. Default: 12.", + ) + args = parser.parse_args() + generate_cuda_fused(args.output_directory, args.num_splits) diff --git a/natten/scripts/autogen_fna.py b/natten/scripts/autogen_fna.py new file mode 100644 index 00000000..09fdcc05 --- /dev/null +++ b/natten/scripts/autogen_fna.py @@ -0,0 +1,906 @@ +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# This script is intended to emit fused kernel instantiations into +# a variable number of source files, generate appropriate headers +# and a single dispatcher interface, which will be used by the +# NATTEN API to call the kernels. +# +# NOTE: these scripts are heavily under-documented, and +# overly-repetitive, and will be replaced in future PRs. +# Please use it with caution. + +import argparse +import os +from typing import List, Tuple + +DEFAULT_OUTPUT_DIR = "csrc/" + +_SM_TO_UPPER_BOUND = { + 50: 70, + 70: 75, + 75: 80, + 80: None, +} + +_SM_TO_LOWER_BOUND = { + 50: 50, + 70: 70, + 75: 75, + 80: 80, +} + + +class KernelConfig: + def __init__(self, na_dim: int, sm: int, gemm_shape: Tuple[int, int, int]): + assert 0 < na_dim <= 3 + assert 50 <= sm <= 80 + assert len(gemm_shape) == 3 + self.na_dim = na_dim + self.sm = sm + self.gemm_shape = gemm_shape + + def get_name(self, is_backward: bool) -> str: + backward_str = "" if not is_backward else "_backward" + return f"fna{self.na_dim}d{backward_str}_{self.gemm_shape[0]}x{self.gemm_shape[1]}x{self.gemm_shape[2]}_sm{self.sm}" + + +class KernelConfigList: + def __init__(self, na_dim: int, sm: int, gemm_shapes: List[Tuple[int, int, int]]): + assert 0 < na_dim <= 3 + assert 50 <= sm <= 80 + self.na_dim = na_dim + self.sm = sm + self.gemm_shapes = gemm_shapes + + @property + def configs(self) -> List[KernelConfig]: + return [ + KernelConfig(na_dim=self.na_dim, sm=self.sm, gemm_shape=gemm_shape) + for gemm_shape in self.gemm_shapes + ] + + def get_name(self, is_backward: bool) -> str: + backward_str = "" if not is_backward else "_backward" + return f"fna{self.na_dim}d{backward_str}_sm{self.sm}" + + +class DataType: + def __init__(self, name, torch_name, short_name, bits, min_sm): + self.name = name + self.torch_name = torch_name + self.bits = bits + self.short_name = short_name + self.min_sm = min_sm + + +NATTEN_Float = DataType("float", "torch::kFloat32", "float32", 32, 0) +NATTEN_Half = DataType("cutlass::half_t", "torch::kFloat16", "float16", 16, 50) +NATTEN_BFloat = DataType("cutlass::bfloat16_t", "torch::kBFloat16", "bfloat16", 16, 80) + + +KERNEL_DECL_TEMPLATE = """__global__ void __launch_bounds__( + {CPP_CLASS}::kNumThreads, + {CPP_CLASS}::kMinBlocksPerSm) +{NAME}(typename {CPP_CLASS}::Params p); +""" + + +KERNEL_IMPL_TEMPLATE = """__global__ void __launch_bounds__( + {CPP_CLASS}::kNumThreads, + {CPP_CLASS}::kMinBlocksPerSm) +{NAME}(typename {CPP_CLASS}::Params p) {{ +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= {SM}0 +#if __CUDA_ARCH__ < {SM_MAX}0 + if (!p.advance_to_block()) {{ + return; + }} + {CPP_CLASS}::attention_kernel(p); + return; +#endif +#endif + printf( + "FATAL: FNA kernel `{NAME}` was built for SM{SM}, but attempted to launch from SM%d\\n", + int(__CUDA_ARCH__ + 0) / 10); +#endif +}} +""" + +SM_TO_SM_MAX = { + 50: 70, + 70: 75, + 75: 80, + 80: 130, +} + + +class FusedNAKernel: + def __init__( + self, + is_backward: bool, + dtype: DataType, + config: KernelConfig, + causal_mask: List, + ): + self.dtype = dtype + self.config = config + self.causal_mask = causal_mask + self.name_cc = f"{self.config.get_name(is_backward)}_{dtype.short_name}" + assert len(self.causal_mask) == self.config.na_dim + cm_str = "_".join([str(int(x)) for x in self.causal_mask]) + self.name_cc += f"_cm_{cm_str}" + self.path_to_header = ( + "natten/cuda/fna/kernel_forward.h" + if not is_backward + else "natten/cuda/fna/kernel_backward.h" + ) + self.aligned = True + self.is_backward = is_backward + + @property + def causal_mask_inst(self) -> str: + return ( + "CausalMask<" + + ", ".join(["true" if x else "false" for x in self.causal_mask]) + + ">" + ) + + @property + def cpp_class(self) -> str: + if self.is_backward: + template_args = ", ".join( + [ + str(self.config.na_dim), + self.causal_mask_inst, + self.dtype.name, + f"cutlass::arch::Sm{self.config.sm}", + "true" if self.aligned else "false", + str(self.config.gemm_shape[0]), + str(self.config.gemm_shape[1]), + str(self.config.gemm_shape[2]), + ] + ) + return f"FusedNeighborhoodAttentionBackwardKernel<{template_args}>" + else: + template_args = ", ".join( + [ + str(self.config.na_dim), + self.causal_mask_inst, + self.dtype.name, + f"cutlass::arch::Sm{self.config.sm}", + "true" if self.aligned else "false", + str(self.config.gemm_shape[0]), + str(self.config.gemm_shape[1]), + str(self.config.gemm_shape[2]), + ] + ) + return f"FusedNeighborhoodAttentionKernel<{template_args}>" + + def header(self): + return KERNEL_DECL_TEMPLATE.format( + CPP_CLASS=self.cpp_class, + NAME=self.name_cc, + ) + + def source(self): + return KERNEL_IMPL_TEMPLATE.format( + CPP_CLASS=self.cpp_class, + NAME=self.name_cc, + SM=self.config.sm, + SM_MAX=SM_TO_SM_MAX[self.config.sm], + ) + + +class FusedNAKernelBundle: + def __init__( + self, + is_backward: bool, + na_dim: int, + sm: int, + dtype: DataType, + config_list: KernelConfigList, + causal_mask: List, + ): + self.na_dim = na_dim + self.sm = sm + self.dtype = dtype + self.config_list = config_list + self.causal_mask = causal_mask + self.name_cc = f"{self.config_list.get_name(is_backward)}_{dtype.short_name}" + assert len(self.causal_mask) == self.config_list.na_dim + cm_str = "_".join([str(int(x)) for x in self.causal_mask]) + self.name_cc += f"_cm_{cm_str}" + self.path_to_header = ( + "natten/cuda/fna/kernel_forward.h" + if not is_backward + else "natten/cuda/fna/kernel_backward.h" + ) + self.aligned = True + self.is_backward = is_backward + + @property + def kernels(self): + return [ + FusedNAKernel( + is_backward=self.is_backward, + dtype=self.dtype, + causal_mask=self.causal_mask, + config=config, + ) + for config in self.config_list.configs + ] + + def header(self): + decl = "\n\n" + decl += "///////////////////////////////////////////////////////////////////\n" + decl += f"// FNA-{self.na_dim}D / {self.dtype.short_name} / SM{self.sm}" + decl += "Backward Kernel" if self.is_backward else "" + decl += "\n" + decl += "///////////////////////////////////////////////////////////////////" + + for kernel in self.kernels: + decl += "\n\n" + decl += kernel.header() + + decl += f"\n\ntemplate \nvoid {self.name_cc}" + decl += "(T cb) {\n" + for kernel in self.kernels: + decl += f" cb({kernel.cpp_class}(), {kernel.name_cc});\n" + decl += "}" + return decl + + def source(self): + impl = "\n\n" + impl += "///////////////////////////////////////////////////////////////////\n" + impl += f"// FNA-{self.na_dim}D / {self.dtype.short_name} / SM{self.sm}\n" + impl += "///////////////////////////////////////////////////////////////////" + + for kernel in self.kernels: + impl += "\n\n" + impl += kernel.source() + return impl + + +def write_combined_source_file(path, filename, headers, sources): + source_head = [] + source_head += ["#include \n"] + source_head += ["#include \n"] + + for header in headers: + source_head += [f"#include <{header}>\n"] + + source_head += ["namespace natten { \n"] + source_head += ["namespace cuda { \n"] + source_head += ["namespace fna { \n\n"] + + source_head = "".join(source_head) + + source_body = [] + for source in sources: + source_body += source.source() + source_body = "".join(source_body) + + source_foot = "".join( + [ + "} \n", + "} \n", + "} \n", + "\n", + ] + ) + filename = f"{path}/{filename}" + with open(filename, "w") as f: + f.write(source_head) + f.write(source_body) + f.write(source_foot) + + +class RankDispatcher: + def __init__(self, is_backward: bool): + self.name_cc = ( + "DISPATCH_FNA_FORWARD_KERNEL" + if not is_backward + else "DISPATCH_FNA_BACKWARD_KERNEL" + ) + self.name_target = ( + "DISPATCH_FNA_FORWARD_" if not is_backward else "DISPATCH_FNA_BACKWARD_" + ) + self.dims: List[int] = [] + self.is_backward = is_backward + + def append(self, na_dim: int): + self.dims.append(na_dim) + + def get_dispatcher(self): + dispatcher_str = "" + if self.is_backward: + dispatcher_str += ( + f"#define {self.name_cc}(rank, cc, dtype, is_causal, cb) \\\n" + ) + else: + dispatcher_str += ( + f"#define {self.name_cc}(rank, cc, dtype, is_causal, cb) \\\n" + ) + dispatcher_str += " [&] { \\\n" + for i, na_dim in enumerate(self.dims): + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + dispatcher_str += f"if constexpr (rank == {na_dim})" + dispatcher_str += " { \\\n" + dispatcher_str += " " + if self.is_backward: + dispatcher_str += ( + f" {self.name_target}{na_dim}D(cc, dtype, is_causal, cb); \\\n" + ) + else: + dispatcher_str += ( + f" {self.name_target}{na_dim}D(cc, dtype, is_causal, cb); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += ( + ' throw std::runtime_error("NATTEN FNA kernel dispatch failed! ' + + "Invalid spatial extent rank! Only 1, 2, and 3D are supported!" + + '"); \\\n' + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class DeviceDispatcher: + def __init__(self, is_backward: bool, na_dim: int): + self.na_dim = na_dim + self.name_cc = ( + f"DISPATCH_FNA_FORWARD_{self.na_dim}D" + if not is_backward + else f"DISPATCH_FNA_BACKWARD_{self.na_dim}D" + ) + self.name_target = self.name_cc + "_SM" + self.devices: List[int] = [] + self.is_backward = is_backward + + def append(self, sm: int): + self.devices.append(sm) + + def get_dispatcher(self): + dispatcher_str = "" + if self.is_backward: + dispatcher_str += f"#define {self.name_cc}(cc, dtype, is_causal, cb) \\\n" + else: + dispatcher_str += f"#define {self.name_cc}(cc, dtype, is_causal, cb) \\\n" + dispatcher_str += " [&] { \\\n" + for i, sm in enumerate(self.devices): + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + dispatcher_str += "if (" + if _SM_TO_UPPER_BOUND[sm] is not None: + dispatcher_str += f"cc < {_SM_TO_UPPER_BOUND[sm]} && " + dispatcher_str += f"cc >= {_SM_TO_LOWER_BOUND[sm]})" + dispatcher_str += " { \\\n" + dispatcher_str += " " + if self.is_backward: + dispatcher_str += ( + f" {self.name_target}{sm}(dtype, is_causal, cb); \\\n" + ) + else: + dispatcher_str += ( + f" {self.name_target}{sm}(dtype, is_causal, cb); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += ( + ' throw std::runtime_error("NATTEN FNA kernel dispatch failed! ' + + "Fused neighborhood attention is not implemented for this device." + + '"); \\\n' + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class DataTypeDispatcher: + def __init__(self, is_backward: bool, na_dim: int, sm: int): + self.dtypes: List[DataType] = [] + self.na_dim = na_dim + self.sm = sm + self.name_cc = ( + f"DISPATCH_FNA_FORWARD_{self.na_dim}D_SM{self.sm}" + if not is_backward + else f"DISPATCH_FNA_BACKWARD_{self.na_dim}D_SM{self.sm}" + ) + self.name_target = self.name_cc + self.is_backward = is_backward + + def append(self, dtype: DataType): + self.dtypes.append(dtype) + + def get_dispatcher(self): + dispatcher_str = "" + if self.is_backward: + dispatcher_str += f"#define {self.name_cc}(dtype, is_causal, cb) \\\n" + else: + dispatcher_str += f"#define {self.name_cc}(dtype, is_causal, cb) \\\n" + dispatcher_str += " [&] { \\\n" + for i, dtype in enumerate(self.dtypes): + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + dispatcher_str += f"if (dtype == {dtype.torch_name})" + dispatcher_str += " { \\\n" + dispatcher_str += " " + if self.is_backward: + dispatcher_str += ( + f" {self.name_target}_{dtype.short_name}(is_causal, cb); \\\n" + ) + else: + dispatcher_str += ( + f" {self.name_target}_{dtype.short_name}(is_causal, cb); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += ( + ' throw std::runtime_error("NATTEN kernel dispatch failed! ' + + f"FNA-{self.na_dim}D does not support this data type on SM{self.sm}." + + '"); \\\n' + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class CausalMaskDispatcher: + def __init__( + self, + is_backward: bool, + na_dim: int, + sm: int, + dtype: DataType, + ): + self.na_dim = na_dim + self.sm = sm + self.dtype = dtype + self.name_cc = ( + (f"DISPATCH_FNA_FORWARD_{self.na_dim}D_SM{self.sm}_{self.dtype.short_name}") + if not is_backward + else ( + f"DISPATCH_FNA_BACKWARD_{self.na_dim}D_SM{self.sm}_{self.dtype.short_name}" + ) + ) + backward_str = "" if not is_backward else "_backward" + self.is_backward = is_backward + self.name_target = ( + f"fna{self.na_dim}d{backward_str}_sm{self.sm}_{self.dtype.short_name}" + ) + self.cms: List = [] + + def append(self, cm): + assert len(cm) == self.na_dim + self.cms.append(cm) + + def get_dispatcher(self): + dispatcher_str = "" + if self.is_backward: + dispatcher_str += f"#define {self.name_cc}(is_causal, cb) \\\n" + else: + dispatcher_str += f"#define {self.name_cc}(is_causal, cb) \\\n" + dispatcher_str += " [&] { \\\n" + i = 0 + for cm in self.cms: + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + i += 1 + cm_str = "_".join([str(int(x)) for x in cm]) + dispatcher_str += "if (" + for dim in range(self.na_dim): + dispatcher_str += ( + f"std::get<{dim}>(is_causal)" + if cm[dim] + else f"!std::get<{dim}>(is_causal)" + ) + if dim != self.na_dim - 1: + dispatcher_str += " && " + dispatcher_str += ")" + dispatcher_str += " { \\\n" + dispatcher_str += " " + dispatcher_str += f" {self.name_target}_cm_{cm_str}(cb); \\\n" + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + maybe_backward_str = "-backward" if self.is_backward else "" + dispatcher_str += ( + ' throw std::runtime_error("NATTEN FNA kernel dispatch failed! ' + + f"Causal mask dispatcher (FNA-{self.na_dim}D{maybe_backward_str}, SM{self.sm}, {self.dtype.short_name}) got invalid causal mask!" + + '"); \\\n' + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +def write_header_file(content, path, namespaces, extra_includes=None): + extra_includes = extra_includes or [] + header_head = [ + "#pragma once\n", + "\n\n", + ] + header_head += ["#include \n"] + header_head += ["#include \n"] + header_head += ["#include \n"] + for incl in extra_includes: + header_head += [f"#include <{incl}> \n"] + + for namespace in namespaces: + header_head += [f"namespace {namespace}", " { \n"] + + header_foot = [ + "\n\n", + ] + for namespace in namespaces: + header_foot += ["} ", f"// namespace {namespace}", " \n"] + header_foot += [ + "\n", + ] + with open(path, "w") as f: + f.write("".join(header_head)) + f.write(content) + f.write("".join(header_foot)) + + +def generate_cuda_kernels(path, num_splits=2): + + SUPPORTED_ARCHS = [50, 70, 75, 80] + + CUDA_DTYPES = [ + NATTEN_Float, + NATTEN_Half, + NATTEN_BFloat, + ] + NA_RANKS = [1, 2, 3] + + CAUSAL_MASKS = { + 1: [[0], [1]], + 2: [[0, 0], [0, 1], [1, 0], [1, 1]], + 3: [ + [0, 0, 0], + [0, 0, 1], + [0, 1, 0], + [0, 1, 1], + [1, 0, 0], + [1, 0, 1], + [1, 1, 0], + [1, 1, 1], + ], + } + + GEMM_SHAPES = [ + # K = 32 + (32, 128, 32), + (64, 64, 32), + (64, 128, 32), + # K = 64 + (32, 128, 64), + (64, 64, 64), + (64, 128, 64), + # K = 128 + (32, 128, 128), + (64, 64, 128), + (64, 128, 128), + # K = 65536 + (32, 128, 2**16), + (64, 128, 2**16), + (64, 64, 2**16), + ] + device_dispatchers = [] + dtype_dispatchers = [] + cm_dispatchers = [] + kernels = [] + + rank_dispatcher = RankDispatcher(is_backward=False) + for na_dim in NA_RANKS: + rank_dispatcher.append(na_dim) + device_dispatcher = DeviceDispatcher(is_backward=False, na_dim=na_dim) + for sm in SUPPORTED_ARCHS: + device_dispatcher.append(sm) + dtype_dispatcher = DataTypeDispatcher( + is_backward=False, na_dim=na_dim, sm=sm + ) + for dtype in CUDA_DTYPES: + if dtype.min_sm > sm: + continue + dtype_dispatcher.append(dtype) + cm_dispatcher = CausalMaskDispatcher( + is_backward=False, na_dim=na_dim, sm=sm, dtype=dtype + ) + for cm in CAUSAL_MASKS[na_dim]: + cm_dispatcher.append(cm) + kernels.append( + FusedNAKernelBundle( + is_backward=False, + na_dim=na_dim, + sm=sm, + dtype=dtype, + config_list=KernelConfigList( + na_dim=na_dim, + sm=sm, + gemm_shapes=GEMM_SHAPES, # GEMM_SHAPES[na_dim][sm] + ), + causal_mask=cm, + ) + ) + cm_dispatchers.append(cm_dispatcher) + dtype_dispatchers.append(dtype_dispatcher) + device_dispatchers.append(device_dispatcher) + + # + + BACKWARD_GEMM_SHAPES = { + 80: { + NATTEN_Half: [ + # K = 32 + (64, 64, 32), + # K = 64 + (64, 64, 64), + # K = 128 + (64, 64, 128), + (128, 128, 128), # SM80/SM90 only + # K = 65536 + (64, 64, 2**16), + (128, 64, 2**16), # SM80/SM90 only + ], + NATTEN_BFloat: [ + # K = 32 + (64, 64, 32), + # K = 64 + (64, 64, 64), + # K = 128 + (64, 64, 128), + (128, 128, 128), # SM80/SM90 only + # K = 65536 + (64, 64, 2**16), + (128, 64, 2**16), # SM80/SM90 only + ], + NATTEN_Float: [ + # K = 32 + (64, 64, 32), + # K = 64 + (64, 64, 64), + # K = 128 + (64, 64, 128), + (128, 64, 128), # SM80/SM90 only + # K = 65536 + (64, 64, 2**16), + (128, 64, 2**16), # SM80/SM90 only + ], + }, + 75: { + NATTEN_Half: [ + # K = 32 + (64, 64, 32), + # K = 64 + (64, 64, 64), + # K = 128 + (64, 64, 128), + # K = 65536 + (64, 64, 2**16), + ], + NATTEN_Float: [ + # K = 32 + (64, 64, 32), + # K = 64 + (64, 64, 64), + # K = 128 + (64, 64, 128), + # K = 65536 + (64, 64, 2**16), + ], + }, + 70: { + NATTEN_Half: [ + # K = 32 + (64, 64, 32), + # K = 64 + (64, 64, 64), + # K = 128 + (64, 64, 128), + (128, 64, 128), # SM70 only? + # K = 65536 + (64, 64, 2**16), + (128, 64, 2**16), # SM70 only? + ], + NATTEN_Float: [ + # K = 32 + (64, 64, 32), + # K = 64 + (64, 64, 64), + # K = 128 + (64, 64, 128), + # K = 65536 + (64, 64, 2**16), + ], + }, + 50: { + NATTEN_Half: [ + # K = 32 + (64, 64, 32), + # K = 64 + (64, 64, 64), + # K = 128 + (64, 64, 128), + # K = 65536 + (64, 64, 2**16), + ], + NATTEN_Float: [ + # K = 32 + (64, 64, 32), + # K = 64 + (64, 64, 64), + # K = 128 + (64, 64, 128), + # K = 65536 + (64, 64, 2**16), + ], + }, + } + + # Backward kernels + rank_dispatcher_backward = RankDispatcher(is_backward=True) + # for na_dim in [1, 2, 3]: + for na_dim in NA_RANKS: + rank_dispatcher_backward.append(na_dim) + device_dispatcher = DeviceDispatcher(is_backward=True, na_dim=na_dim) + for sm in SUPPORTED_ARCHS: + device_dispatcher.append(sm) + dtype_dispatcher = DataTypeDispatcher( + is_backward=True, na_dim=na_dim, sm=sm + ) + for dtype in CUDA_DTYPES: + if dtype.min_sm > sm: + continue + dtype_dispatcher.append(dtype) + cm_dispatcher = CausalMaskDispatcher( + is_backward=True, na_dim=na_dim, sm=sm, dtype=dtype + ) + for cm in CAUSAL_MASKS[na_dim]: + cm_dispatcher.append(cm) + kernels.append( + FusedNAKernelBundle( + is_backward=True, + na_dim=na_dim, + sm=sm, + dtype=dtype, + config_list=KernelConfigList( + na_dim=na_dim, + sm=sm, + gemm_shapes=BACKWARD_GEMM_SHAPES[sm][ + dtype + ], # GEMM_SHAPES[na_dim][sm] + ), + causal_mask=cm, + ) + ) + cm_dispatchers.append(cm_dispatcher) + dtype_dispatchers.append(dtype_dispatcher) + device_dispatchers.append(device_dispatcher) + + # + + path_to_sources = f"{path}/autogen/src/cuda/fna/" + rel_header = "natten_autogen/cuda/fna/" + path_to_header_dir = f"{path}/autogen/include/{rel_header}" + + os.makedirs(path_to_sources, exist_ok=False) + os.makedirs(path_to_header_dir, exist_ok=False) + + path_headers = f"{path_to_header_dir}kernels.h" + path_rank = f"{path_to_header_dir}interface.h" + path_device = f"{path_to_header_dir}dispatch_device.h" + path_dtype = f"{path_to_header_dir}dispatch_dtype.h" + path_cm = f"{path_to_header_dir}dispatch_cm.h" + + rel_path_headers = f"{rel_header}kernels.h" + rel_path_device = f"{rel_header}dispatch_device.h" + rel_path_dtype = f"{rel_header}dispatch_dtype.h" + rel_path_cm = f"{rel_header}dispatch_cm.h" + + rank_disp = rank_dispatcher.get_dispatcher() + rank_disp += rank_dispatcher_backward.get_dispatcher() + + device_disp = "" + for dispatcher in device_dispatchers: + device_disp += dispatcher.get_dispatcher() + + dtype_disp = "" + for dispatcher in dtype_dispatchers: + dtype_disp += dispatcher.get_dispatcher() + + cm_disp = "" + for dispatcher in cm_dispatchers: + cm_disp += dispatcher.get_dispatcher() + + headers = "" + for kernel in kernels: + headers += kernel.header() + + assert ( + len(kernels) >= num_splits + ), f"Generated {len(kernels)} kernels, but got {num_splits=}." + split_size = len(kernels) // num_splits + num_splits_with_res = len(kernels) % num_splits + kernels_emitted = [] + kernels_split = [] + for split_idx in range(num_splits): + kernel_start_idx = split_size * split_idx + min(num_splits_with_res, split_idx) + num_kernels_in_split = split_size + ( + 1 if split_idx < num_splits_with_res else 0 + ) + kernel_end_idx = kernel_start_idx + num_kernels_in_split + assert kernel_end_idx <= len(kernels) + pth_set = set() + source_list = [] + for kernel_idx in range(kernel_start_idx, kernel_end_idx): + kernel = kernels[kernel_idx] + pth_set.add(kernel.path_to_header) + source_list.append(kernel) + kernels_emitted.append(kernel_idx) + pth_set.add(rel_path_headers) + write_combined_source_file( + path_to_sources, f"source_{split_idx}.cu", sorted(pth_set), source_list + ) + kernels_split.append(source_list) + # print(f"{split_idx=}, {kernel_start_idx=}, {kernel_end_idx=}, {len(kernels_emitted)=}") + assert split_idx == num_splits - 1, f"Expected {split_idx=} == {num_splits=} - 1" + assert len(kernels_emitted) == len(kernels) + assert sorted(kernels_emitted) == [ + x for x in range(len(kernels)) + ], f"{sorted(kernels_emitted)=}" + assert all(len(x) > 0 for x in kernels_split) + + namespaces = ["natten", "cuda", "fna"] + cuda_headers = [ + "natten/natten.h", + "natten/cuda/fna/na_utils.cuh", + "natten/cuda/fna/kernel_forward.h", + "natten/cuda/fna/kernel_backward.h", + ] + write_header_file( + rank_disp, path_rank, namespaces, cuda_headers + [rel_path_device] + ) + write_header_file( + device_disp, path_device, namespaces, cuda_headers + [rel_path_dtype] + ) + write_header_file(dtype_disp, path_dtype, namespaces, cuda_headers + [rel_path_cm]) + write_header_file(cm_disp, path_cm, namespaces, cuda_headers + [rel_path_headers]) + # write_header_file(ks_disp, path_ks, namespaces, cuda_headers + [rel_path_di]) + # write_header_file(di_disp, path_di, namespaces, cuda_headers + [rel_path_headers]) + write_header_file(headers, path_headers, namespaces, cuda_headers) + + +def generate_cuda_fused(output_directory: str, num_splits: int): + generate_cuda_kernels(output_directory, num_splits=num_splits) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "-o", + "--output-directory", + default=DEFAULT_OUTPUT_DIR, + help="Path to the directory where the auto-generated " + "kernel instantiations are dumped. " + f"Default: {DEFAULT_OUTPUT_DIR}", + ) + parser.add_argument( + "--num-splits", + type=int, + default=128, + help="Number of source files into which the kernels are split. Default: 128.", + ) + args = parser.parse_args() + generate_cuda_fused(args.output_directory, args.num_splits) diff --git a/natten/scripts/autogen_hopper_fmha.py b/natten/scripts/autogen_hopper_fmha.py new file mode 100644 index 00000000..29769c81 --- /dev/null +++ b/natten/scripts/autogen_hopper_fmha.py @@ -0,0 +1,669 @@ +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# This script is intended to emit fused kernel instantiations into +# a variable number of source files, generate appropriate headers +# and a single dispatcher interface, which will be used by the +# NATTEN API to call the kernels. +# +# NOTE: these scripts are heavily under-documented, and +# overly-repetitive, and will be replaced in future PRs. +# Please use it with caution. + +import argparse +import os +from enum import Enum +from typing import List, Tuple + +DEFAULT_OUTPUT_DIR = "csrc/" + + +class KernelType(Enum): + NonPersistent = 0 + WSCooperative = 1 + WSPingpong = 2 + + +SUPPORTED_CONFIGS_FORWARD = { + 16: { + 32: [ + ((64, 128, 32), KernelType.NonPersistent), + ((128, 64, 32), KernelType.WSCooperative), + ], + 64: [ + ((64, 128, 64), KernelType.NonPersistent), + ((128, 64, 64), KernelType.WSCooperative), + ((128, 64, 64), KernelType.WSPingpong), + ], + 128: [ + ((128, 128, 128), KernelType.WSCooperative), + ((128, 128, 128), KernelType.WSPingpong), + ], + 256: [ + ((128, 64, 256), KernelType.WSCooperative), + ], + }, + # TODO: FP8 + 8: { + 32: [], + 64: [], + 128: [], + 256: [], + }, +} + + +KERNEL_DECL_TEMPLATE = """ +void {kernel_name}( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); +""" + + +KERNEL_IMPL_TEMPLATE = """ +void {kernel_name}( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) {{ + + using GemmShape = {GEMMShape}; + using Kernel = natten::cuda::fmha_hopper::KernelForward< + {dtype}, GemmShape, {kernel_type}, cutlass::fmha::collective::DefaultFusion, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_hopper::KernelForward< + {dtype}, GemmShape, {kernel_type}, cutlass::fmha::collective::ResidualFusion, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_hopper::KernelForward< + {dtype}, GemmShape, {kernel_type}, cutlass::fmha::collective::CausalFusion, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_hopper::KernelForward< + {dtype}, + GemmShape, + {kernel_type}, + cutlass::fmha::collective::ResidualFusion, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_hopper::KernelForward< + {dtype}, + GemmShape, + {kernel_type}, + cutlass::fmha::collective::CausalFusion, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) {{ + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({{bytes}}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }}; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_k % get<1>(GemmShape{{}}) == 0; + if (no_mask_required) {{ + Kernel kernel; + launch_kernel(kernel); + }} + else if (is_varlen && is_causal) {{ + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + }} + else if (is_varlen && not is_causal) {{ + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + }} + else if (not is_varlen && is_causal) {{ + KernelWithCausalMask kernel; + launch_kernel(kernel); + }} + else {{ + KernelWithResidualMask kernel; + launch_kernel(kernel); + }} +}} +""" + + +class DataType: + def __init__(self, name, short_name, torch_name, bits): + self.name = name + self.bits = bits + self.short_name = short_name + self.torch_name = torch_name + + +Half = DataType("cutlass::half_t", "float16", "torch::kFloat16", 16) +BFloat = DataType("cutlass::bfloat16_t", "bfloat16", "torch::kBFloat16", 16) + + +def iterable_to_static_cute_tuple(shape_in) -> str: + shape = ", ".join([f"cute::Int<{x}>" for x in shape_in]) + return f"cute::tuple<{shape}>" + + +def kernel_type_to_str(kernel_type: KernelType) -> str: + namespace = "natten::cuda::hopper::HopperKernelSchedule::" + if kernel_type == KernelType.NonPersistent: + return namespace + "NonPersistent" + if kernel_type == KernelType.WSCooperative: + return namespace + "WSCooperative" + if kernel_type == KernelType.WSPingpong: + return namespace + "WSPingpong" + + raise NotImplementedError() + + +def kernel_type_to_tag(kernel_type: KernelType) -> str: + if kernel_type == KernelType.NonPersistent: + return "" + if kernel_type == KernelType.WSCooperative: + return "_coop" + if kernel_type == KernelType.WSPingpong: + return "_pp" + + raise NotImplementedError() + + +class HopperFmhaInstance: + def __init__( + self, + dtype: DataType, + gemm_shape: Tuple[int, int, int], + kernel_type: KernelType, + ): + assert len(gemm_shape) == 3 + self.gemm_shape = gemm_shape + self.dtype = dtype + self.max_head_dim = gemm_shape[2] + self.kernel_type = kernel_type + + def get_gemm_shape_cute(self) -> str: + return iterable_to_static_cute_tuple(self.gemm_shape) + + def get_name(self) -> str: + name = "hopper_fmha" + name += f"_{self.dtype.short_name}" + name += "_" + "x".join([str(x) for x in self.gemm_shape]) + name += kernel_type_to_tag(self.kernel_type) + return name + + def get_decl(self) -> str: + return KERNEL_DECL_TEMPLATE.format( + kernel_name=self.get_name(), + ) + + def get_impl(self) -> str: + return KERNEL_IMPL_TEMPLATE.format( + kernel_name=self.get_name(), + GEMMShape=self.get_gemm_shape_cute(), + kernel_type=kernel_type_to_str(self.kernel_type), + dtype=self.dtype.name, + ) + + +def write_combined_source_file(path, filename, headers, kernels): + source_head = [] + source_head += ["#ifdef NATTEN_WITH_CUTLASS\n"] + source_head += ["#ifdef NATTEN_WITH_HOPPER_FNA\n"] + + source_head += ["#include \n"] + source_head += ["#include \n"] + + source_head += ["#include \n"] + source_head += ["#include \n"] + source_head += ["#include \n"] + source_head += ["#include \n"] + + source_head += ["#include \n"] + source_head += ["#include \n"] + + source_head += ["#include \n"] + + for header in headers: + source_head += [f"#include <{header}>\n"] + + source_head += ["namespace natten { \n"] + source_head += ["namespace cuda { \n"] + source_head += ["namespace fmha_hopper { \n\n"] + + source_head = "".join(source_head) + + source_body = [] + for kernel in kernels: + source_body += "\n\n" + kernel.get_impl() + "\n\n" + source_body = "".join(source_body) + + source_foot = "".join( + [ + "} // namespace fmha_hopper \n", + "} // namespace cuda \n", + "} // namespace natten \n", + "#endif \n", + "#endif \n", + "\n", + ] + ) + filename = f"{path}/{filename}" + with open(filename, "w") as f: + f.write(source_head) + f.write(source_body) + f.write(source_foot) + + +class DTypeDispatcher: + def __init__(self): + self.dtypes = [] + self.name = "DISPATCH_HOPPER_FMHA_FORWARD" + + def append(self, dtype: DataType): + self.dtypes.append(dtype) + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += f"#define {self.name}(dtype, dim, q_tile_size, kv_tile_size, kernel_type, ...) \\\n" + dispatcher_str += " [&] { \\\n" + for i, dtype in enumerate(self.dtypes): + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + dispatcher_str += f"if (dtype == {dtype.torch_name})" + dispatcher_str += " { \\\n" + dispatcher_str += " " + dispatcher_str += f" {self.name}_{dtype.short_name}(dim, q_tile_size, kv_tile_size, kernel_type, __VA_ARGS__); \\\n" + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += ( + ' throw std::runtime_error("Hopper FMHA forward kernel dispatch failed! ' + + 'It does not support dtype " + std::string(c10::toString(dtype)) + "."' + + "); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class HeadDimDispatcher: + def __init__(self, dtype: DataType): + self.dtype = dtype + self.name = f"DISPATCH_HOPPER_FMHA_FORWARD_{self.dtype.short_name}" + + self.dims: List[int] = [] + + def append(self, dim: int): + self.dims.append(dim) + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += f"#define {self.name}(dim, q_tile_size, kv_tile_size, kernel_type, ...) \\\n" + dispatcher_str += " [&] { \\\n" + for i, dim in enumerate(self.dims): + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + dispatcher_str += f"if (dim == {dim})" + dispatcher_str += " { \\\n" + dispatcher_str += " " + dispatcher_str += f" {self.name}_headdim{dim}(q_tile_size, kv_tile_size, kernel_type, __VA_ARGS__); \\\n" + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += ( + " throw std::runtime_error(" + '"Hopper FMHA forward kernel dispatch failed! ' + 'It does not support head dim "' + + f' + std::to_string(dim) + " for {self.dtype.short_name}."' + + "); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class ConfigDispatcher: + def __init__( + self, + dtype: DataType, + head_dim: int, + ): + self.dtype = dtype + self.head_dim = head_dim + + self.name = ( + f"DISPATCH_HOPPER_FMHA_FORWARD_{self.dtype.short_name}_headdim{head_dim}" + ) + + self.configs: List = [] + + def append(self, config): + assert len(config) == 3 + assert isinstance(config[0], int) + assert isinstance(config[1], int) + assert isinstance(config[2], KernelType) + self.configs.append(config) + + def get_kernel_instance(self, q_tile_size, kv_tile_size, kernel_type): + gemm_M = q_tile_size + gemm_N = kv_tile_size + gemm_K = self.head_dim + gemm_shape = (gemm_M, gemm_N, gemm_K) + config = (gemm_shape, kernel_type) + + supported_configs = SUPPORTED_CONFIGS_FORWARD[self.dtype.bits][self.head_dim] + assert ( + config in supported_configs + ), f"{config=} not in supported configs {supported_configs=}" + + kernel = HopperFmhaInstance( + dtype=self.dtype, + gemm_shape=gemm_shape, + kernel_type=kernel_type, + ) + + return kernel + + def get_target_name(self, q_tile_size, kv_tile_size, kernel_type): + kernel = self.get_kernel_instance(q_tile_size, kv_tile_size, kernel_type) + return kernel.get_name() + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += ( + f"#define {self.name}(q_tile_size, kv_tile_size, kernel_type, ...) \\\n" + ) + dispatcher_str += " [&] { \\\n" + i = 0 + for q_tile_size, kv_tile_size, kernel_type in self.configs: + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + i += 1 + dispatcher_str += "if (" + dispatcher_str += f"q_tile_size == {q_tile_size}" + dispatcher_str += " && " + dispatcher_str += "\\\n" + dispatcher_str += f"kv_tile_size == {kv_tile_size}" + dispatcher_str += " && " + dispatcher_str += "\\\n" + dispatcher_str += f"kernel_type == {kernel_type_to_str(kernel_type)}" + dispatcher_str += ")" + dispatcher_str += " { \\\n" + + dispatcher_str += f" natten::cuda::fmha_hopper::{self.get_target_name(q_tile_size, kv_tile_size, kernel_type)}(__VA_ARGS__); \\\n" + + dispatcher_str += "} \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += " " + dispatcher_str += ( + " throw std::runtime_error(" + '"Hopper FMHA forward kernel dispatch failed! ' + "It got invalid Q tile, KV tile, and schedule " + + f"combination ({self.dtype.short_name}, " + f"head_dim {self.head_dim}): " + 'q_tile=" + std::to_string(q_tile_size)' + ' + ", kv_tile=" + std::to_string(kv_tile_size)' + ' + ", schedule="' + " + natten::cuda::hopper::to_string(kernel_type)" + ' + "."' + "); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +def write_header_file(content, path, namespaces, extra_includes=None): + extra_includes = extra_includes or [] + header_head = [ + "#pragma once\n", + "\n\n", + ] + header_head += ["#include \n"] + header_head += ["#include \n"] + header_head += ["#include \n"] + header_head += ["#include \n"] + header_head += ["#ifdef NATTEN_WITH_CUTLASS\n"] + header_head += ["#ifdef NATTEN_WITH_HOPPER_FNA\n"] + for incl in extra_includes: + header_head += [f"#include <{incl}> \n"] + + for namespace in namespaces: + header_head += [f"namespace {namespace}", " { \n"] + + header_foot = [ + "\n\n", + ] + for namespace in namespaces: + header_foot += ["} ", f"// namespace {namespace}", " \n"] + header_foot += [ + "#endif \n", + "#endif \n", + "\n", + ] + with open(path, "w") as f: + f.write("".join(header_head)) + f.write(content) + f.write("".join(header_foot)) + + +def generate_hopper_fmha_kernels(path, num_splits=2): + + SUPPORTED_DTYPES = [ + Half, + BFloat, + ] + + HEAD_DIMS = [32, 64, 128, 256] + + CONFIGS = { + 16: { + 32: [ + (64, 128, KernelType.NonPersistent), + ], + 64: [ + (64, 128, KernelType.NonPersistent), + ], + 128: [ + (128, 128, KernelType.WSCooperative), + (128, 128, KernelType.WSPingpong), + ], + 256: [ + (128, 64, KernelType.WSCooperative), + ], + }, + } + + head_dim_dispatchers = [] + config_dispatchers = [] + kernels = [] + + dtype_dispatcher = DTypeDispatcher() + for dtype in SUPPORTED_DTYPES: + dtype_dispatcher.append(dtype) + + head_dim_dispatcher = HeadDimDispatcher(dtype=dtype) + + for head_dim in HEAD_DIMS: + head_dim_dispatcher.append(head_dim) + + config_dispatcher = ConfigDispatcher( + dtype=dtype, + head_dim=head_dim, + ) + + for q_tile_size, kv_tile_size, kernel_type in CONFIGS[dtype.bits][head_dim]: + config_dispatcher.append((q_tile_size, kv_tile_size, kernel_type)) + kernels.append( + config_dispatcher.get_kernel_instance( + q_tile_size, kv_tile_size, kernel_type + ) + ) + config_dispatchers.append(config_dispatcher) + head_dim_dispatchers.append(head_dim_dispatcher) + + # + + path_to_sources = f"{path}/autogen/src/cuda/hopper_fmha/" + rel_header = "natten_autogen/cuda/hopper_fmha/" + path_to_header_dir = f"{path}/autogen/include/{rel_header}" + + os.makedirs(path_to_sources, exist_ok=False) + os.makedirs(path_to_header_dir, exist_ok=False) + + path_headers = f"{path_to_header_dir}kernels.h" + path_dtype = f"{path_to_header_dir}interface.h" + path_head_dim = f"{path_to_header_dir}dispatch_head_dim.h" + path_tile_size = f"{path_to_header_dir}dispatch_tile_size.h" + + rel_path_headers = f"{rel_header}kernels.h" + rel_path_head_dim = f"{rel_header}dispatch_head_dim.h" + rel_path_tile_size = f"{rel_header}dispatch_tile_size.h" + + dtype_disp = dtype_dispatcher.get_dispatcher() + + head_dim_disp = "" + for dispatcher in head_dim_dispatchers: + head_dim_disp += dispatcher.get_dispatcher() + + config_disp = "" + for dispatcher in config_dispatchers: + config_disp += dispatcher.get_dispatcher() + + headers = "" + for kernel in kernels: + headers += kernel.get_decl() + + assert ( + len(kernels) >= num_splits + ), f"Generated {len(kernels)} kernels, but got {num_splits=}." + split_size = len(kernels) // num_splits + num_splits_with_res = len(kernels) % num_splits + kernels_emitted = [] + kernels_split = [] + for split_idx in range(num_splits): + kernel_start_idx = split_size * split_idx + min(num_splits_with_res, split_idx) + num_kernels_in_split = split_size + ( + 1 if split_idx < num_splits_with_res else 0 + ) + kernel_end_idx = kernel_start_idx + num_kernels_in_split + assert kernel_end_idx <= len(kernels) + pth_set = set() + source_list = [] + for kernel_idx in range(kernel_start_idx, kernel_end_idx): + kernel = kernels[kernel_idx] + # pth_set.add(kernel.path_to_header) + source_list.append(kernel) + kernels_emitted.append(kernel_idx) + pth_set.add(rel_path_headers) + write_combined_source_file( + path_to_sources, f"source_{split_idx}.cu", sorted(pth_set), source_list + ) + kernels_split.append(source_list) + # print(f"{split_idx=}, {kernel_start_idx=}, {kernel_end_idx=}, {len(kernels_emitted)=}") + assert split_idx == num_splits - 1, f"Expected {split_idx=} == {num_splits=} - 1" + assert len(kernels_emitted) == len(kernels) + assert sorted(kernels_emitted) == [ + x for x in range(len(kernels)) + ], f"{sorted(kernels_emitted)=}" + assert all(len(x) > 0 for x in kernels_split) + + namespaces = ["natten", "cuda", "fmha_hopper"] + cuda_headers = [ + "natten/natten.h", + "ATen/ATen.h", + "ATen/cuda/CUDAContext.h", + "c10/cuda/CUDAGuard.h", + "c10/cuda/CUDAStream.h", + "natten/natten.h", + "natten/helpers.h", + "natten/cuda/hopper_fmha_fna.h", + "natten/cuda/fmha_hopper/fmha_forward.cuh", + ] + write_header_file( + dtype_disp, path_dtype, namespaces, cuda_headers + [rel_path_head_dim] + ) + write_header_file( + head_dim_disp, path_head_dim, namespaces, cuda_headers + [rel_path_tile_size] + ) + write_header_file( + config_disp, path_tile_size, namespaces, cuda_headers + [rel_path_headers] + ) + write_header_file(headers, path_headers, namespaces, cuda_headers) + + +def generate_hopper_fmha(output_directory: str, num_splits: int): + generate_hopper_fmha_kernels(output_directory, num_splits=num_splits) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "-o", + "--output-directory", + default=DEFAULT_OUTPUT_DIR, + help="Path to the directory where the auto-generated " + "kernel instantiations are dumped. " + f"Default: {DEFAULT_OUTPUT_DIR}", + ) + parser.add_argument( + "--num-splits", + type=int, + default=1, + help="Number of source files into which the kernels are split. Default: 1.", + ) + args = parser.parse_args() + generate_hopper_fmha(args.output_directory, args.num_splits) diff --git a/natten/scripts/autogen_hopper_fmha_bwd.py b/natten/scripts/autogen_hopper_fmha_bwd.py new file mode 100644 index 00000000..9ecbbb48 --- /dev/null +++ b/natten/scripts/autogen_hopper_fmha_bwd.py @@ -0,0 +1,620 @@ +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# This script is intended to emit fused kernel instantiations into +# a variable number of source files, generate appropriate headers +# and a single dispatcher interface, which will be used by the +# NATTEN API to call the kernels. +# +# NOTE: these scripts are heavily under-documented, and +# overly-repetitive, and will be replaced in future PRs. +# Please use it with caution. + +import argparse +import os +from typing import List, Tuple + +DEFAULT_OUTPUT_DIR = "csrc/" + + +SUPPORTED_CONFIGS_BACKWARD = { + 16: { + 32: [ + (64, 128, 32), + (128, 128, 32), + ], + 64: [ + (64, 128, 64), + (128, 128, 64), + ], + 128: [ + (64, 128, 128), + ], + }, +} + + +KERNEL_DECL_TEMPLATE = """ +void {kernel_name}( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options); +""" + + +KERNEL_IMPL_TEMPLATE = """ +void {kernel_name}( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + bool is_causal, + float attn_scale, + // varlen parameters + bool is_varlen, + int max_seqlen_Q, + int max_seqlen_KV, + void* ptr_cumulative_seqlen_Q, + void* ptr_cumulative_seqlen_KV, + // init/launch params + int device_id, + cudaStream_t stream, + at::TensorOptions tensor_options) {{ + + using GemmShape = {GEMMShape}; + using Kernel = natten::cuda::fmha_hopper::KernelBackward< + {dtype}, GemmShape, cutlass::fmha::collective::DefaultFusion, /* kIsVarlen= */ false>; + using KernelWithResidualMask = natten::cuda::fmha_hopper::KernelBackward< + {dtype}, GemmShape, cutlass::fmha::collective::ResidualFusion, /* kIsVarlen= */ false>; + using KernelWithCausalMask = natten::cuda::fmha_hopper::KernelBackward< + {dtype}, GemmShape, cutlass::fmha::collective::CausalFusion, /* kIsVarlen= */ false>; + + // Varlen kernels + using VarlenKernelWithResidualMask = natten::cuda::fmha_hopper::KernelBackward< + {dtype}, + GemmShape, + cutlass::fmha::collective::ResidualFusion, + /* kIsVarlen= */ true>; + using VarlenKernelWithCausalMask = natten::cuda::fmha_hopper::KernelBackward< + {dtype}, + GemmShape, + cutlass::fmha::collective::CausalFusion, + /* kIsVarlen= */ true>; + + auto launch_kernel = [&](auto& kernel) {{ + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + attn_scale, + // varlen + max_seqlen_Q, + max_seqlen_KV, + ptr_cumulative_seqlen_Q, + ptr_cumulative_seqlen_KV, + // + device_id); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({{bytes}}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); + }}; + + bool no_mask_required = not is_varlen && not is_causal && seqlen_q % get<0>(GemmShape{{}}) == 0 && seqlen_k % get<1>(GemmShape{{}}) == 0; + if (no_mask_required) {{ + Kernel kernel; + launch_kernel(kernel); + }} + else if (is_varlen && is_causal) {{ + VarlenKernelWithCausalMask kernel; + launch_kernel(kernel); + }} + else if (is_varlen && not is_causal) {{ + VarlenKernelWithResidualMask kernel; + launch_kernel(kernel); + }} + else if (not is_varlen && is_causal) {{ + KernelWithCausalMask kernel; + launch_kernel(kernel); + }} + else {{ + KernelWithResidualMask kernel; + launch_kernel(kernel); + }} +}} +""" + + +class DataType: + def __init__(self, name, short_name, torch_name, bits): + self.name = name + self.bits = bits + self.short_name = short_name + self.torch_name = torch_name + + +Half = DataType("cutlass::half_t", "float16", "torch::kFloat16", 16) +BFloat = DataType("cutlass::bfloat16_t", "bfloat16", "torch::kBFloat16", 16) + + +def iterable_to_static_cute_tuple(shape_in) -> str: + shape = ", ".join([f"cute::Int<{x}>" for x in shape_in]) + return f"cute::tuple<{shape}>" + + +class HopperFmhaInstance: + def __init__( + self, + dtype: DataType, + gemm_shape: Tuple[int, int, int], + ): + assert len(gemm_shape) == 3 + self.gemm_shape = gemm_shape + self.dtype = dtype + self.max_head_dim = gemm_shape[2] + + def get_gemm_shape_cute(self) -> str: + return iterable_to_static_cute_tuple(self.gemm_shape) + + def get_name(self) -> str: + name = "hopper_fmha_backward" + name += f"_{self.dtype.short_name}" + name += "_" + "x".join([str(x) for x in self.gemm_shape]) + return name + + def get_decl(self) -> str: + return KERNEL_DECL_TEMPLATE.format( + kernel_name=self.get_name(), + ) + + def get_impl(self) -> str: + return KERNEL_IMPL_TEMPLATE.format( + kernel_name=self.get_name(), + GEMMShape=self.get_gemm_shape_cute(), + dtype=self.dtype.name, + ) + + +def write_combined_source_file(path, filename, headers, kernels): + source_head = [] + source_head += ["#ifdef NATTEN_WITH_CUTLASS\n"] + source_head += ["#ifdef NATTEN_WITH_HOPPER_FNA\n"] + + source_head += ["#include \n"] + source_head += ["#include \n"] + + source_head += ["#include \n"] + source_head += ["#include \n"] + source_head += ["#include \n"] + source_head += ["#include \n"] + + source_head += ["#include \n"] + source_head += ["#include \n"] + + source_head += ["#include \n"] + + for header in headers: + source_head += [f"#include <{header}>\n"] + + source_head += ["namespace natten { \n"] + source_head += ["namespace cuda { \n"] + source_head += ["namespace fmha_hopper { \n\n"] + + source_head = "".join(source_head) + + source_body = [] + for kernel in kernels: + source_body += "\n\n" + kernel.get_impl() + "\n\n" + source_body = "".join(source_body) + + source_foot = "".join( + [ + "} // namespace fmha_hopper \n", + "} // namespace cuda \n", + "} // namespace natten \n", + "#endif \n", + "#endif \n", + "\n", + ] + ) + filename = f"{path}/{filename}" + with open(filename, "w") as f: + f.write(source_head) + f.write(source_body) + f.write(source_foot) + + +class DTypeDispatcher: + def __init__(self): + self.dtypes = [] + self.name = "DISPATCH_HOPPER_FMHA_BACKWARD" + + def append(self, dtype: DataType): + self.dtypes.append(dtype) + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += ( + f"#define {self.name}(dtype, dim, q_tile_size, kv_tile_size, ...) \\\n" + ) + dispatcher_str += " [&] { \\\n" + for i, dtype in enumerate(self.dtypes): + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + dispatcher_str += f"if (dtype == {dtype.torch_name})" + dispatcher_str += " { \\\n" + dispatcher_str += " " + dispatcher_str += f" {self.name}_{dtype.short_name}(dim, q_tile_size, kv_tile_size, __VA_ARGS__); \\\n" + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += ( + ' throw std::runtime_error("Hopper FMHA backward kernel dispatch failed! ' + + 'It does not support dtype " + std::string(c10::toString(dtype)) + "."' + + "); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class HeadDimDispatcher: + def __init__(self, dtype: DataType): + self.dtype = dtype + self.name = f"DISPATCH_HOPPER_FMHA_BACKWARD_{self.dtype.short_name}" + + self.dims: List[int] = [] + + def append(self, dim: int): + self.dims.append(dim) + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += ( + f"#define {self.name}(dim, q_tile_size, kv_tile_size, ...) \\\n" + ) + dispatcher_str += " [&] { \\\n" + for i, dim in enumerate(self.dims): + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + dispatcher_str += f"if (dim == {dim})" + dispatcher_str += " { \\\n" + dispatcher_str += " " + dispatcher_str += f" {self.name}_headdim{dim}(q_tile_size, kv_tile_size, __VA_ARGS__); \\\n" + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += ( + " throw std::runtime_error(" + '"Hopper FMHA backward kernel dispatch failed! ' + 'It does not support head dim "' + + f' + std::to_string(dim) + " for {self.dtype.short_name}."' + + "); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class ConfigDispatcher: + def __init__( + self, + dtype: DataType, + head_dim: int, + ): + self.dtype = dtype + self.head_dim = head_dim + + self.name = ( + f"DISPATCH_HOPPER_FMHA_BACKWARD_{self.dtype.short_name}_headdim{head_dim}" + ) + + self.configs: List = [] + + def append(self, config): + assert len(config) == 2 + assert isinstance(config[0], int) + assert isinstance(config[1], int) + self.configs.append(config) + + def get_kernel_instance(self, q_tile_size, kv_tile_size): + gemm_M = q_tile_size + gemm_N = kv_tile_size + gemm_K = self.head_dim + gemm_shape = (gemm_M, gemm_N, gemm_K) + + supported_configs = SUPPORTED_CONFIGS_BACKWARD[self.dtype.bits][self.head_dim] + assert ( + gemm_shape in supported_configs + ), f"{gemm_shape=} not in supported configs {supported_configs=}" + + kernel = HopperFmhaInstance( + dtype=self.dtype, + gemm_shape=gemm_shape, + ) + + return kernel + + def get_target_name(self, q_tile_size, kv_tile_size): + kernel = self.get_kernel_instance(q_tile_size, kv_tile_size) + return kernel.get_name() + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += f"#define {self.name}(q_tile_size, kv_tile_size, ...) \\\n" + dispatcher_str += " [&] { \\\n" + i = 0 + for q_tile_size, kv_tile_size in self.configs: + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + i += 1 + dispatcher_str += "if (" + dispatcher_str += f"q_tile_size == {q_tile_size}" + dispatcher_str += " && " + dispatcher_str += "\\\n" + dispatcher_str += f"kv_tile_size == {kv_tile_size})" + dispatcher_str += " { \\\n" + + dispatcher_str += f" natten::cuda::fmha_hopper::{self.get_target_name(q_tile_size, kv_tile_size)}(__VA_ARGS__); \\\n" + + dispatcher_str += "} \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += " " + dispatcher_str += ( + " throw std::runtime_error(" + '"Hopper FMHA backward kernel dispatch failed! ' + "It got invalid Q tile and KV tile combination " + + f"({self.dtype.short_name}, head_dim {self.head_dim}): " + 'q_tile=" + std::to_string(q_tile_size)' + ' + ", kv_tile=" + std::to_string(kv_tile_size)' + ' + "."' + "); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +def write_header_file(content, path, namespaces, extra_includes=None): + extra_includes = extra_includes or [] + header_head = [ + "#pragma once\n", + "\n\n", + ] + header_head += ["#include \n"] + header_head += ["#include \n"] + header_head += ["#include \n"] + header_head += ["#include \n"] + header_head += ["#ifdef NATTEN_WITH_CUTLASS\n"] + header_head += ["#ifdef NATTEN_WITH_HOPPER_FNA\n"] + for incl in extra_includes: + header_head += [f"#include <{incl}> \n"] + + for namespace in namespaces: + header_head += [f"namespace {namespace}", " { \n"] + + header_foot = [ + "\n\n", + ] + for namespace in namespaces: + header_foot += ["} ", f"// namespace {namespace}", " \n"] + header_foot += [ + "#endif \n", + "#endif \n", + "\n", + ] + with open(path, "w") as f: + f.write("".join(header_head)) + f.write(content) + f.write("".join(header_foot)) + + +def generate_hopper_fmha_kernels(path, num_splits=2): + + SUPPORTED_DTYPES = [ + Half, + BFloat, + ] + + HEAD_DIMS = [32, 64, 128] + + CONFIGS = { + 16: { + 32: [ + (64, 128), + (128, 128), + ], + 64: [ + (64, 128), + (128, 128), + ], + 128: [ + (64, 128), + ], + }, + } + + head_dim_dispatchers = [] + config_dispatchers = [] + kernels = [] + + dtype_dispatcher = DTypeDispatcher() + for dtype in SUPPORTED_DTYPES: + dtype_dispatcher.append(dtype) + + head_dim_dispatcher = HeadDimDispatcher(dtype=dtype) + + for head_dim in HEAD_DIMS: + head_dim_dispatcher.append(head_dim) + + config_dispatcher = ConfigDispatcher( + dtype=dtype, + head_dim=head_dim, + ) + + for q_tile_size, kv_tile_size in CONFIGS[dtype.bits][head_dim]: + config_dispatcher.append((q_tile_size, kv_tile_size)) + kernels.append( + config_dispatcher.get_kernel_instance(q_tile_size, kv_tile_size) + ) + config_dispatchers.append(config_dispatcher) + head_dim_dispatchers.append(head_dim_dispatcher) + + # + + path_to_sources = f"{path}/autogen/src/cuda/hopper_fmha_bwd/" + rel_header = "natten_autogen/cuda/hopper_fmha_bwd/" + path_to_header_dir = f"{path}/autogen/include/{rel_header}" + + os.makedirs(path_to_sources, exist_ok=False) + os.makedirs(path_to_header_dir, exist_ok=False) + + path_headers = f"{path_to_header_dir}kernels.h" + path_dtype = f"{path_to_header_dir}interface.h" + path_head_dim = f"{path_to_header_dir}dispatch_head_dim.h" + path_tile_size = f"{path_to_header_dir}dispatch_tile_size.h" + + rel_path_headers = f"{rel_header}kernels.h" + rel_path_head_dim = f"{rel_header}dispatch_head_dim.h" + rel_path_tile_size = f"{rel_header}dispatch_tile_size.h" + + dtype_disp = dtype_dispatcher.get_dispatcher() + + head_dim_disp = "" + for dispatcher in head_dim_dispatchers: + head_dim_disp += dispatcher.get_dispatcher() + + config_disp = "" + for dispatcher in config_dispatchers: + config_disp += dispatcher.get_dispatcher() + + headers = "" + for kernel in kernels: + headers += kernel.get_decl() + + assert ( + len(kernels) >= num_splits + ), f"Generated {len(kernels)} kernels, but got {num_splits=}." + split_size = len(kernels) // num_splits + num_splits_with_res = len(kernels) % num_splits + kernels_emitted = [] + kernels_split = [] + for split_idx in range(num_splits): + kernel_start_idx = split_size * split_idx + min(num_splits_with_res, split_idx) + num_kernels_in_split = split_size + ( + 1 if split_idx < num_splits_with_res else 0 + ) + kernel_end_idx = kernel_start_idx + num_kernels_in_split + assert kernel_end_idx <= len(kernels) + pth_set = set() + source_list = [] + for kernel_idx in range(kernel_start_idx, kernel_end_idx): + kernel = kernels[kernel_idx] + # pth_set.add(kernel.path_to_header) + source_list.append(kernel) + kernels_emitted.append(kernel_idx) + pth_set.add(rel_path_headers) + write_combined_source_file( + path_to_sources, f"source_{split_idx}.cu", sorted(pth_set), source_list + ) + kernels_split.append(source_list) + # print(f"{split_idx=}, {kernel_start_idx=}, {kernel_end_idx=}, {len(kernels_emitted)=}") + assert split_idx == num_splits - 1, f"Expected {split_idx=} == {num_splits=} - 1" + assert len(kernels_emitted) == len(kernels) + assert sorted(kernels_emitted) == [ + x for x in range(len(kernels)) + ], f"{sorted(kernels_emitted)=}" + assert all(len(x) > 0 for x in kernels_split) + + namespaces = ["natten", "cuda", "fmha_hopper"] + cuda_headers = [ + "natten/natten.h", + "ATen/ATen.h", + "ATen/cuda/CUDAContext.h", + "c10/cuda/CUDAGuard.h", + "c10/cuda/CUDAStream.h", + "natten/natten.h", + "natten/helpers.h", + "natten/cuda/fmha_hopper/fmha_backward.cuh", + ] + write_header_file( + dtype_disp, path_dtype, namespaces, cuda_headers + [rel_path_head_dim] + ) + write_header_file( + head_dim_disp, path_head_dim, namespaces, cuda_headers + [rel_path_tile_size] + ) + write_header_file( + config_disp, path_tile_size, namespaces, cuda_headers + [rel_path_headers] + ) + write_header_file(headers, path_headers, namespaces, cuda_headers) + + +def generate_hopper_fmha(output_directory: str, num_splits: int): + generate_hopper_fmha_kernels(output_directory, num_splits=num_splits) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "-o", + "--output-directory", + default=DEFAULT_OUTPUT_DIR, + help="Path to the directory where the auto-generated " + "kernel instantiations are dumped. " + f"Default: {DEFAULT_OUTPUT_DIR}", + ) + parser.add_argument( + "--num-splits", + type=int, + default=1, + help="Number of source files into which the kernels are split. Default: 1.", + ) + args = parser.parse_args() + generate_hopper_fmha(args.output_directory, args.num_splits) diff --git a/natten/scripts/autogen_hopper_fna.py b/natten/scripts/autogen_hopper_fna.py new file mode 100644 index 00000000..e7504a43 --- /dev/null +++ b/natten/scripts/autogen_hopper_fna.py @@ -0,0 +1,885 @@ +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# This script is intended to emit fused kernel instantiations into +# a variable number of source files, generate appropriate headers +# and a single dispatcher interface, which will be used by the +# NATTEN API to call the kernels. +# +# NOTE: these scripts are heavily under-documented, and +# overly-repetitive, and will be replaced in future PRs. +# Please use it with caution. + +import argparse +import math +import os +from enum import Enum +from typing import List, Tuple + +DEFAULT_OUTPUT_DIR = "csrc/" + + +class KernelType(Enum): + NonPersistent = 0 + WSCooperative = 1 + WSPingpong = 2 + + +SUPPORTED_CONFIGS_FORWARD = { + 16: { + 32: [ + ((64, 128, 32), KernelType.NonPersistent), + ((128, 64, 32), KernelType.WSCooperative), + ], + 64: [ + ((64, 128, 64), KernelType.NonPersistent), + ((128, 64, 64), KernelType.WSCooperative), + ((128, 64, 64), KernelType.WSPingpong), + ], + 128: [ + ((128, 128, 128), KernelType.WSCooperative), + ((128, 128, 128), KernelType.WSPingpong), + ], + 256: [ + ((128, 64, 256), KernelType.WSCooperative), + ], + }, + # TODO: FP8 + 8: { + 32: [], + 64: [], + 128: [], + 256: [], + }, +} + + +KERNEL_DECL_TEMPLATE = """ +void {kernel_name}( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + {DimType} q_shape, + {DimType} kv_shape, + {DimType} qkv_shape, + {DimType} window_size, + {DimType} stride, + {DimType} dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); +""" + + +KERNEL_IMPL_TEMPLATE = """ +void {kernel_name}( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + {DimType} q_shape, + {DimType} kv_shape, + {DimType} qkv_shape, + {DimType} window_size, + {DimType} stride, + {DimType} dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) {{ + + using Causal = {Causal}; + using QTileShape = {QTileShape}; + using KVTileShape = {KVTileShape}; + using GemmShape = {GEMMShape}; + using Kernel = natten::cuda::fna_hopper::KernelForward< + {dtype}, Causal, QTileShape, KVTileShape, GemmShape, {kernel_type}>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({{bytes}}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +}} +""" + + +class DataType: + def __init__(self, name, short_name, torch_name, bits): + self.name = name + self.bits = bits + self.short_name = short_name + self.torch_name = torch_name + + +Half = DataType("cutlass::half_t", "float16", "torch::kFloat16", 16) +BFloat = DataType("cutlass::bfloat16_t", "bfloat16", "torch::kBFloat16", 16) + + +def iterable_to_static_cute_tuple(shape_in) -> str: + shape = ", ".join([f"cute::Int<{x}>" for x in shape_in]) + return f"cute::tuple<{shape}>" + + +def get_dim_type(na_dim: int) -> str: + shape = ", ".join(["int" for _ in range(na_dim)]) + return f"cute::tuple<{shape}>" + + +def kernel_type_to_str(kernel_type: KernelType) -> str: + namespace = "natten::cuda::hopper::HopperKernelSchedule::" + if kernel_type == KernelType.NonPersistent: + return namespace + "NonPersistent" + if kernel_type == KernelType.WSCooperative: + return namespace + "WSCooperative" + if kernel_type == KernelType.WSPingpong: + return namespace + "WSPingpong" + + raise NotImplementedError() + + +def kernel_type_to_tag(kernel_type: KernelType) -> str: + if kernel_type == KernelType.NonPersistent: + return "" + if kernel_type == KernelType.WSCooperative: + return "_coop" + if kernel_type == KernelType.WSPingpong: + return "_pp" + + raise NotImplementedError() + + +class HopperFnaInstance: + def __init__( + self, + na_dim: int, + dtype: DataType, + gemm_shape: Tuple[int, int, int], + kernel_type: KernelType, + q_tile_shape: tuple, + kv_tile_shape: tuple, + causal: tuple, + ): + assert 0 < na_dim <= 3 + assert len(gemm_shape) == 3 + assert na_dim == len(q_tile_shape) + assert na_dim == len(kv_tile_shape) + assert na_dim == len(causal) + self.na_dim = na_dim + self.gemm_shape = gemm_shape + self.q_tile_shape = q_tile_shape + self.kv_tile_shape = kv_tile_shape + self.causal = causal + self.dtype = dtype + self.max_head_dim = gemm_shape[2] + self.kernel_type = kernel_type + + def get_q_tile_shape_cute(self) -> str: + return iterable_to_static_cute_tuple(self.q_tile_shape) + + def get_kv_tile_shape_cute(self) -> str: + return iterable_to_static_cute_tuple(self.kv_tile_shape) + + def get_causal_cute(self) -> str: + consts = ", ".join( + ["cute::true_type" if c else "cute::false_type" for c in self.causal] + ) + return f"cute::tuple<{consts}>" + + def get_gemm_shape_cute(self) -> str: + return iterable_to_static_cute_tuple(self.gemm_shape) + + def get_name(self) -> str: + name = f"hopper_fna{self.na_dim}d" + name += f"_{self.dtype.short_name}" + name += "_" + "x".join([str(x) for x in self.gemm_shape]) + name += kernel_type_to_tag(self.kernel_type) + name += "_Q" + "x".join([str(x) for x in self.q_tile_shape]) + name += "_KV" + "x".join([str(x) for x in self.kv_tile_shape]) + name += "_causal" + "x".join(["1" if c else "0" for c in self.causal]) + return name + + def get_decl(self) -> str: + return KERNEL_DECL_TEMPLATE.format( + kernel_name=self.get_name(), DimType=get_dim_type(self.na_dim) + ) + + def get_impl(self) -> str: + return KERNEL_IMPL_TEMPLATE.format( + kernel_name=self.get_name(), + DimType=get_dim_type(self.na_dim), + Causal=self.get_causal_cute(), + QTileShape=self.get_q_tile_shape_cute(), + KVTileShape=self.get_kv_tile_shape_cute(), + GEMMShape=self.get_gemm_shape_cute(), + kernel_type=kernel_type_to_str(self.kernel_type), + dtype=self.dtype.name, + ) + + +def write_combined_source_file(path, filename, headers, kernels): + source_head = [] + source_head += ["#ifdef NATTEN_WITH_CUTLASS\n"] + source_head += ["#ifdef NATTEN_WITH_HOPPER_FNA\n"] + + source_head += ["#include \n"] + source_head += ["#include \n"] + + source_head += ["#include \n"] + source_head += ["#include \n"] + source_head += ["#include \n"] + source_head += ["#include \n"] + + source_head += ["#include \n"] + source_head += ["#include \n"] + + source_head += ["#include \n"] + + for header in headers: + source_head += [f"#include <{header}>\n"] + + source_head += ["namespace natten { \n"] + source_head += ["namespace cuda { \n"] + source_head += ["namespace fna_hopper { \n\n"] + + source_head = "".join(source_head) + + source_body = [] + for kernel in kernels: + source_body += "\n\n" + kernel.get_impl() + "\n\n" + source_body = "".join(source_body) + + source_foot = "".join( + [ + "} // namespace fna_hopper \n", + "} // namespace cuda \n", + "} // namespace natten \n", + "#endif \n", + "#endif \n", + "\n", + ] + ) + filename = f"{path}/{filename}" + with open(filename, "w") as f: + f.write(source_head) + f.write(source_body) + f.write(source_foot) + + +class NaDimDispatcher: + def __init__( + self, + ): + self.name = "DISPATCH_HOPPER_FNA_FORWARD" + self.dims = [] + + def append(self, na_dim: int): + self.dims.append(na_dim) + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += f"#define {self.name}(rank, dtype, dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \\\n" + dispatcher_str += " [&] { \\\n" + for i, na_dim in enumerate(self.dims): + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + dispatcher_str += f"if constexpr (rank == {na_dim})" + dispatcher_str += " { \\\n" + dispatcher_str += " " + dispatcher_str += f" {self.name}_{na_dim}D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \\\n" + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += ( + ' throw std::runtime_error("Hopper FNA forward kernel dispatch failed! ' + + "It only supports NA1D, 2D, and 3D!" + + '"); \\\n' + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class DTypeDispatcher: + def __init__(self, na_dim: int): + self.dtypes: List[DataType] = [] + self.na_dim = na_dim + self.name = f"DISPATCH_HOPPER_FNA_FORWARD_{self.na_dim}D" + + def append(self, dtype: DataType): + self.dtypes.append(dtype) + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += f"#define {self.name}(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \\\n" + dispatcher_str += " [&] { \\\n" + for i, dtype in enumerate(self.dtypes): + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + dispatcher_str += f"if (dtype == {dtype.torch_name})" + dispatcher_str += " { \\\n" + dispatcher_str += " " + dispatcher_str += f" {self.name}_{dtype.short_name}(dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \\\n" + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += ( + ' throw std::runtime_error("Hopper FNA-' + + f"{self.na_dim}D forward kernel dispatch failed! " + 'It does not support dtype "' + ' + std::string(c10::toString(dtype)) + "."' + "); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class HeadDimDispatcher: + def __init__(self, na_dim: int, dtype: DataType): + self.na_dim = na_dim + self.dtype = dtype + self.name = ( + f"DISPATCH_HOPPER_FNA_FORWARD_{self.na_dim}D_{self.dtype.short_name}" + ) + + self.dims: List[int] = [] + + def append(self, dim: int): + self.dims.append(dim) + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += f"#define {self.name}(dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \\\n" + dispatcher_str += " [&] { \\\n" + for i, dim in enumerate(self.dims): + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + dispatcher_str += f"if (dim == {dim})" + dispatcher_str += " { \\\n" + dispatcher_str += " " + dispatcher_str += f" {self.name}_headdim{dim}(is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \\\n" + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += ( + " throw std::runtime_error(" + '"Hopper FNA forward kernel dispatch failed! ' + 'It does not support head dim "' + + f' + std::to_string(dim) + " for {self.dtype.short_name}."' + + "); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class CausalMaskDispatcher: + def __init__( + self, + na_dim: int, + dtype: DataType, + head_dim: int, + ): + self.na_dim = na_dim + self.dtype = dtype + self.head_dim = head_dim + + self.name = f"DISPATCH_HOPPER_FNA_FORWARD_{self.na_dim}D_{self.dtype.short_name}_headdim{head_dim}" + + self.cms: List = [] + + def append(self, cm): + assert len(cm) == self.na_dim, f"{cm} incompatible for {self.na_dim}D NA." + self.cms.append(cm) + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += f"#define {self.name}(is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \\\n" + dispatcher_str += " [&] { \\\n" + i = 0 + for cm in self.cms: + cm_str = "causal" + "x".join(["1" if c else "0" for c in cm]) + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + i += 1 + dispatcher_str += "if (" + for dim in range(self.na_dim): + dispatcher_str += ( + f"cute::get<{dim}>(is_causal)" + if cm[dim] + else f"not cute::get<{dim}>(is_causal)" + ) + if dim != self.na_dim - 1: + dispatcher_str += " && " + dispatcher_str += ")" + dispatcher_str += " { \\\n" + dispatcher_str += " " + dispatcher_str += f" {self.name}_{cm_str}(q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \\\n" + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += " " + dispatcher_str += ( + ' throw std::runtime_error("Hopper FNA-' + + f"{self.na_dim}D forward kernel dispatch failed! " + f"Causal mask dispatcher ({self.dtype.short_name}, " + f'head_dim {self.head_dim}) got invalid causal mask!"' + "); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class ConfigDispatcher: + def __init__( + self, + na_dim: int, + dtype: DataType, + head_dim: int, + causal_mask: tuple, + ): + self.na_dim = na_dim + self.dtype = dtype + self.head_dim = head_dim + self.causal_mask = causal_mask + cm_str = "_causal" + "x".join(["1" if c else "0" for c in self.causal_mask]) + + self.name = f"DISPATCH_HOPPER_FNA_FORWARD_{self.na_dim}D_{self.dtype.short_name}_headdim{head_dim}" + self.name += cm_str + + self.configs: List = [] + + def append(self, config): + assert len(config) == 3 + assert isinstance(config[0], tuple) and len(config[0]) == self.na_dim + assert isinstance(config[1], tuple) and len(config[1]) == self.na_dim + assert isinstance(config[2], KernelType) + self.configs.append(config) + + def get_kernel_instance(self, q_tile_shape, kv_tile_shape, kernel_type): + gemm_M = math.prod(q_tile_shape) + gemm_N = math.prod(kv_tile_shape) + gemm_K = self.head_dim + gemm_shape = (gemm_M, gemm_N, gemm_K) + config = (gemm_shape, kernel_type) + + supported_configs = SUPPORTED_CONFIGS_FORWARD[self.dtype.bits][self.head_dim] + assert ( + config in supported_configs + ), f"{config=} not in supported configs {supported_configs=}" + + kernel = HopperFnaInstance( + na_dim=self.na_dim, + dtype=self.dtype, + gemm_shape=gemm_shape, + kernel_type=kernel_type, + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + causal=self.causal_mask, + ) + + return kernel + + def get_target_name(self, q_tile_shape, kv_tile_shape, kernel_type): + kernel = self.get_kernel_instance(q_tile_shape, kv_tile_shape, kernel_type) + return kernel.get_name() + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += ( + f"#define {self.name}(q_tile_shape, kv_tile_shape, kernel_type, ...) \\\n" + ) + dispatcher_str += " [&] { \\\n" + i = 0 + for q_tile_shape, kv_tile_shape, kernel_type in self.configs: + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + i += 1 + dispatcher_str += "if (" + for dim in range(self.na_dim): + dispatcher_str += ( + f"cute::get<{dim}>(q_tile_shape) == {q_tile_shape[dim]}" + ) + dispatcher_str += " && " + dispatcher_str += "\\\n" + for dim in range(self.na_dim): + dispatcher_str += ( + f"cute::get<{dim}>(kv_tile_shape) == {kv_tile_shape[dim]}" + ) + dispatcher_str += " && " + dispatcher_str += "\\\n" + dispatcher_str += f"kernel_type == {kernel_type_to_str(kernel_type)}" + dispatcher_str += ")" + dispatcher_str += " { \\\n" + + dispatcher_str += f" natten::cuda::fna_hopper::{self.get_target_name(q_tile_shape, kv_tile_shape, kernel_type)}(__VA_ARGS__); \\\n" + + dispatcher_str += "} \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += " " + q_tile_str = ' + "," + '.join( + [ + f"std::to_string(cute::get<{i}>(q_tile_shape))" + for i in range(self.na_dim) + ] + ) + kv_tile_str = ' + "," + '.join( + [ + f"std::to_string(cute::get<{i}>(kv_tile_shape))" + for i in range(self.na_dim) + ] + ) + dispatcher_str += ( + ' throw std::runtime_error("Hopper FNA-' + + f"{self.na_dim}D forward kernel dispatch failed! " + "It got invalid Q tile, KV tile, and schedule " + + f"combination ({self.dtype.short_name}, " + f"head_dim {self.head_dim}): " + 'q_tile=(" + ' + q_tile_str + ' + "), ' + 'kv_tile=(" + ' + kv_tile_str + ' + "), ' + 'schedule="' + " + natten::cuda::hopper::to_string(kernel_type)" + ' + "."' + "); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +def write_header_file(content, path, namespaces, extra_includes=None): + extra_includes = extra_includes or [] + header_head = [ + "#pragma once\n", + "\n\n", + ] + header_head += ["#include \n"] + header_head += ["#include \n"] + header_head += ["#include \n"] + header_head += ["#include \n"] + header_head += ["#ifdef NATTEN_WITH_CUTLASS\n"] + header_head += ["#ifdef NATTEN_WITH_HOPPER_FNA\n"] + for incl in extra_includes: + header_head += [f"#include <{incl}> \n"] + + for namespace in namespaces: + header_head += [f"namespace {namespace}", " { \n"] + + header_foot = [ + "\n\n", + ] + for namespace in namespaces: + header_foot += ["} ", f"// namespace {namespace}", " \n"] + header_foot += [ + "#endif \n", + "#endif \n", + "\n", + ] + with open(path, "w") as f: + f.write("".join(header_head)) + f.write(content) + f.write("".join(header_foot)) + + +def generate_hopper_fna_kernels(path, num_splits=2): + + NA_DIMS = [1, 2, 3] + + SUPPORTED_DTYPES = [ + Half, + BFloat, + ] + + HEAD_DIMS = [32, 64, 128, 256] + + CAUSAL_MASKS = { + 1: [(False,), (True,)], + 2: [(False, False), (False, True), (True, False), (True, True)], + 3: [ + (False, False, False), + (False, False, True), + (False, True, False), + (False, True, True), + (True, False, False), + (True, False, True), + (True, True, False), + (True, True, True), + ], + } + + CONFIGS = { + 1: { + 16: { + 32: [ + ((64,), (128,), KernelType.NonPersistent), + ], + 64: [ + ((64,), (128,), KernelType.NonPersistent), + ], + 128: [ + ((128,), (128,), KernelType.WSCooperative), + ((128,), (128,), KernelType.WSPingpong), + ], + 256: [ + ((128,), (64,), KernelType.WSCooperative), + ], + }, + }, + 2: { + 16: { + 32: [ + ((8, 8), (16, 8), KernelType.NonPersistent), + ((8, 8), (8, 16), KernelType.NonPersistent), + ], + 64: [ + ((8, 8), (16, 8), KernelType.NonPersistent), + ((8, 8), (8, 16), KernelType.NonPersistent), + ], + 128: [ + ((16, 8), (16, 8), KernelType.WSCooperative), + ((16, 8), (16, 8), KernelType.WSPingpong), + ], + 256: [ + ((16, 8), (8, 8), KernelType.WSCooperative), + ((8, 16), (8, 8), KernelType.WSCooperative), + ], + }, + }, + 3: { + 16: { + 32: [ + ((4, 4, 4), (4, 4, 8), KernelType.NonPersistent), + ((4, 4, 4), (2, 8, 8), KernelType.NonPersistent), + ], + 64: [ + ((4, 4, 4), (4, 4, 8), KernelType.NonPersistent), + ((4, 4, 4), (2, 8, 8), KernelType.NonPersistent), + ], + 128: [ + ((4, 4, 8), (4, 4, 8), KernelType.WSCooperative), + ((4, 4, 8), (4, 4, 8), KernelType.WSPingpong), + ((2, 8, 8), (2, 8, 8), KernelType.WSCooperative), + ((2, 8, 8), (2, 8, 8), KernelType.WSPingpong), + ], + 256: [ + ((4, 4, 8), (4, 4, 4), KernelType.WSCooperative), + ((2, 8, 8), (4, 4, 4), KernelType.WSCooperative), + ], + }, + }, + } + + dtype_dispatchers = [] + head_dim_dispatchers = [] + cm_dispatchers = [] + config_dispatchers = [] + kernels = [] + + rank_dispatcher = NaDimDispatcher() + for na_dim in NA_DIMS: + rank_dispatcher.append(na_dim) + + dtype_dispatcher = DTypeDispatcher(na_dim=na_dim) + for dtype in SUPPORTED_DTYPES: + dtype_dispatcher.append(dtype) + + head_dim_dispatcher = HeadDimDispatcher(dtype=dtype, na_dim=na_dim) + + for head_dim in HEAD_DIMS: + head_dim_dispatcher.append(head_dim) + + cm_dispatcher = CausalMaskDispatcher( + na_dim=na_dim, dtype=dtype, head_dim=head_dim + ) + for cm in CAUSAL_MASKS[na_dim]: + cm_dispatcher.append(cm) + + config_dispatcher = ConfigDispatcher( + dtype=dtype, + head_dim=head_dim, + na_dim=na_dim, + causal_mask=cm, + ) + + for q_tile_shape, kv_tile_shape, kernel_type in CONFIGS[na_dim][ + dtype.bits + ][head_dim]: + config_dispatcher.append( + (q_tile_shape, kv_tile_shape, kernel_type) + ) + kernels.append( + config_dispatcher.get_kernel_instance( + q_tile_shape, kv_tile_shape, kernel_type + ) + ) + config_dispatchers.append(config_dispatcher) + cm_dispatchers.append(cm_dispatcher) + head_dim_dispatchers.append(head_dim_dispatcher) + dtype_dispatchers.append(dtype_dispatcher) + + # + + path_to_sources = f"{path}/autogen/src/cuda/hopper_fna/" + rel_header = "natten_autogen/cuda/hopper_fna/" + path_to_header_dir = f"{path}/autogen/include/{rel_header}" + + os.makedirs(path_to_sources, exist_ok=False) + os.makedirs(path_to_header_dir, exist_ok=False) + + path_headers = f"{path_to_header_dir}kernels.h" + path_rank = f"{path_to_header_dir}interface.h" + path_dtype = f"{path_to_header_dir}dispatch_dtype.h" + path_head_dim = f"{path_to_header_dir}dispatch_head_dim.h" + path_cm = f"{path_to_header_dir}dispatch_cm.h" + path_tile_shape = f"{path_to_header_dir}dispatch_tile_shape.h" + + rel_path_headers = f"{rel_header}kernels.h" + rel_path_dtype = f"{rel_header}dispatch_dtype.h" + rel_path_head_dim = f"{rel_header}dispatch_head_dim.h" + rel_path_cm = f"{rel_header}dispatch_cm.h" + rel_path_tile_shape = f"{rel_header}dispatch_tile_shape.h" + + rank_disp = rank_dispatcher.get_dispatcher() + + dtype_disp = "" + for dispatcher in dtype_dispatchers: + dtype_disp += dispatcher.get_dispatcher() + + head_dim_disp = "" + for dispatcher in head_dim_dispatchers: + head_dim_disp += dispatcher.get_dispatcher() + + cm_disp = "" + for dispatcher in cm_dispatchers: + cm_disp += dispatcher.get_dispatcher() + + config_disp = "" + for dispatcher in config_dispatchers: + config_disp += dispatcher.get_dispatcher() + + headers = "" + for kernel in kernels: + headers += kernel.get_decl() + + assert ( + len(kernels) >= num_splits + ), f"Generated {len(kernels)} kernels, but got {num_splits=}." + split_size = len(kernels) // num_splits + num_splits_with_res = len(kernels) % num_splits + kernels_emitted = [] + kernels_split = [] + for split_idx in range(num_splits): + kernel_start_idx = split_size * split_idx + min(num_splits_with_res, split_idx) + num_kernels_in_split = split_size + ( + 1 if split_idx < num_splits_with_res else 0 + ) + kernel_end_idx = kernel_start_idx + num_kernels_in_split + assert kernel_end_idx <= len(kernels) + pth_set = set() + source_list = [] + for kernel_idx in range(kernel_start_idx, kernel_end_idx): + kernel = kernels[kernel_idx] + # pth_set.add(kernel.path_to_header) + source_list.append(kernel) + kernels_emitted.append(kernel_idx) + pth_set.add(rel_path_headers) + write_combined_source_file( + path_to_sources, f"source_{split_idx}.cu", sorted(pth_set), source_list + ) + kernels_split.append(source_list) + # print(f"{split_idx=}, {kernel_start_idx=}, {kernel_end_idx=}, {len(kernels_emitted)=}") + assert split_idx == num_splits - 1, f"Expected {split_idx=} == {num_splits=} - 1" + assert len(kernels_emitted) == len(kernels) + assert sorted(kernels_emitted) == [ + x for x in range(len(kernels)) + ], f"{sorted(kernels_emitted)=}" + assert all(len(x) > 0 for x in kernels_split) + + namespaces = ["natten", "cuda", "fna_hopper"] + cuda_headers = [ + "natten/natten.h", + "ATen/ATen.h", + "ATen/cuda/CUDAContext.h", + "c10/cuda/CUDAGuard.h", + "c10/cuda/CUDAStream.h", + "natten/natten.h", + "natten/helpers.h", + "natten/cuda/hopper_fmha_fna.h", + "natten/cuda/fna_hopper/fna_forward.cuh", + ] + write_header_file(rank_disp, path_rank, namespaces, cuda_headers + [rel_path_dtype]) + write_header_file( + dtype_disp, path_dtype, namespaces, cuda_headers + [rel_path_head_dim] + ) + write_header_file( + head_dim_disp, path_head_dim, namespaces, cuda_headers + [rel_path_cm] + ) + write_header_file( + cm_disp, path_cm, namespaces, cuda_headers + [rel_path_tile_shape] + ) + write_header_file( + config_disp, path_tile_shape, namespaces, cuda_headers + [rel_path_headers] + ) + write_header_file(headers, path_headers, namespaces, cuda_headers) + + +def generate_hopper_fna(output_directory: str, num_splits: int): + generate_hopper_fna_kernels(output_directory, num_splits=num_splits) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "-o", + "--output-directory", + default=DEFAULT_OUTPUT_DIR, + help="Path to the directory where the auto-generated " + "kernel instantiations are dumped. " + f"Default: {DEFAULT_OUTPUT_DIR}", + ) + parser.add_argument( + "--num-splits", + type=int, + default=16, + help="Number of source files into which the kernels are split. Default: 16.", + ) + args = parser.parse_args() + generate_hopper_fna(args.output_directory, args.num_splits) diff --git a/natten/scripts/autogen_hopper_fna_bwd.py b/natten/scripts/autogen_hopper_fna_bwd.py new file mode 100644 index 00000000..d3d4770d --- /dev/null +++ b/natten/scripts/autogen_hopper_fna_bwd.py @@ -0,0 +1,838 @@ +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# This script is intended to emit fused kernel instantiations into +# a variable number of source files, generate appropriate headers +# and a single dispatcher interface, which will be used by the +# NATTEN API to call the kernels. +# +# NOTE: these scripts are heavily under-documented, and +# overly-repetitive, and will be replaced in future PRs. +# Please use it with caution. + +import argparse +import math +import os +from typing import List, Tuple + +DEFAULT_OUTPUT_DIR = "csrc/" + +SUPPORTED_CONFIGS_BACKWARD = { + 16: { + 32: [ + (64, 128, 32), + (128, 128, 32), + ], + 64: [ + (64, 128, 64), + (128, 128, 64), + ], + 128: [ + (64, 128, 128), + ], + }, +} + + +KERNEL_DECL_TEMPLATE = """ +void {kernel_name}( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + {DimType} q_shape, + {DimType} kv_shape, + {DimType} qkv_shape, + {DimType} window_size, + {DimType} stride, + {DimType} dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options); +""" + + +KERNEL_IMPL_TEMPLATE = """ +void {kernel_name}( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + void* ptr_dQ, + void* ptr_dK, + void* ptr_dV, + void* ptr_dO, + int batch_size, + int seqlen_q, + int seqlen_k, + int heads, + int dim, + {DimType} q_shape, + {DimType} kv_shape, + {DimType} qkv_shape, + {DimType} window_size, + {DimType} stride, + {DimType} dilation, + int device_id, + float attn_scale, + cudaStream_t stream, + at::TensorOptions tensor_options) {{ + + using Causal = {Causal}; + using QTileShape = {QTileShape}; + using KVTileShape = {KVTileShape}; + using GemmShape = {GEMMShape}; + using Kernel = natten::cuda::fna_hopper::KernelBackward< + {dtype}, Causal, QTileShape, KVTileShape, GemmShape>; + + Kernel kernel; + auto args = kernel.initialize( + ptr_Q, + ptr_K, + ptr_V, + ptr_O, + ptr_LSE, + ptr_dQ, + ptr_dK, + ptr_dV, + ptr_dO, + batch_size, + seqlen_q, + seqlen_k, + heads, + dim, + q_shape, + kv_shape, + qkv_shape, + window_size, + stride, + dilation, + device_id, + attn_scale); + + auto bytes = static_cast(kernel.get_workspace_size(args)); + auto workspace = at::empty({{bytes}}, tensor_options.dtype(at::ScalarType::Byte)); + auto workspace_ptr = static_cast(workspace.data_ptr()); + kernel.run(args, workspace_ptr, stream); +}} +""" + + +class DataType: + def __init__(self, name, short_name, torch_name, bits): + self.name = name + self.bits = bits + self.short_name = short_name + self.torch_name = torch_name + + +Half = DataType("cutlass::half_t", "float16", "torch::kFloat16", 16) +BFloat = DataType("cutlass::bfloat16_t", "bfloat16", "torch::kBFloat16", 16) + + +def iterable_to_static_cute_tuple(shape_in) -> str: + shape = ", ".join([f"cute::Int<{x}>" for x in shape_in]) + return f"cute::tuple<{shape}>" + + +def get_dim_type(na_dim: int) -> str: + shape = ", ".join(["int" for _ in range(na_dim)]) + return f"cute::tuple<{shape}>" + + +class HopperFnaInstance: + def __init__( + self, + na_dim: int, + dtype: DataType, + gemm_shape: Tuple[int, int, int], + q_tile_shape: tuple, + kv_tile_shape: tuple, + causal: tuple, + ): + assert 0 < na_dim <= 3 + assert len(gemm_shape) == 3 + assert na_dim == len(q_tile_shape) + assert na_dim == len(kv_tile_shape) + assert na_dim == len(causal) + self.na_dim = na_dim + self.gemm_shape = gemm_shape + self.q_tile_shape = q_tile_shape + self.kv_tile_shape = kv_tile_shape + self.causal = causal + self.dtype = dtype + self.max_head_dim = gemm_shape[2] + + def get_q_tile_shape_cute(self) -> str: + return iterable_to_static_cute_tuple(self.q_tile_shape) + + def get_kv_tile_shape_cute(self) -> str: + return iterable_to_static_cute_tuple(self.kv_tile_shape) + + def get_causal_cute(self) -> str: + consts = ", ".join( + ["cute::true_type" if c else "cute::false_type" for c in self.causal] + ) + return f"cute::tuple<{consts}>" + + def get_gemm_shape_cute(self) -> str: + return iterable_to_static_cute_tuple(self.gemm_shape) + + def get_name(self) -> str: + name = f"hopper_fna{self.na_dim}d_backward" + name += f"_{self.dtype.short_name}" + name += "_" + "x".join([str(x) for x in self.gemm_shape]) + name += "_Q" + "x".join([str(x) for x in self.q_tile_shape]) + name += "_KV" + "x".join([str(x) for x in self.kv_tile_shape]) + name += "_causal" + "x".join(["1" if c else "0" for c in self.causal]) + return name + + def get_decl(self) -> str: + return KERNEL_DECL_TEMPLATE.format( + kernel_name=self.get_name(), DimType=get_dim_type(self.na_dim) + ) + + def get_impl(self) -> str: + return KERNEL_IMPL_TEMPLATE.format( + kernel_name=self.get_name(), + DimType=get_dim_type(self.na_dim), + Causal=self.get_causal_cute(), + QTileShape=self.get_q_tile_shape_cute(), + KVTileShape=self.get_kv_tile_shape_cute(), + GEMMShape=self.get_gemm_shape_cute(), + dtype=self.dtype.name, + ) + + +def write_combined_source_file(path, filename, headers, kernels): + source_head = [] + source_head += ["#ifdef NATTEN_WITH_CUTLASS\n"] + source_head += ["#ifdef NATTEN_WITH_HOPPER_FNA\n"] + + source_head += ["#include \n"] + source_head += ["#include \n"] + + source_head += ["#include \n"] + source_head += ["#include \n"] + source_head += ["#include \n"] + source_head += ["#include \n"] + + source_head += ["#include \n"] + source_head += ["#include \n"] + + source_head += ["#include \n"] + + for header in headers: + source_head += [f"#include <{header}>\n"] + + source_head += ["namespace natten { \n"] + source_head += ["namespace cuda { \n"] + source_head += ["namespace fna_hopper { \n\n"] + + source_head = "".join(source_head) + + source_body = [] + for kernel in kernels: + source_body += "\n\n" + kernel.get_impl() + "\n\n" + source_body = "".join(source_body) + + source_foot = "".join( + [ + "} // namespace fna_hopper \n", + "} // namespace cuda \n", + "} // namespace natten \n", + "#endif \n", + "#endif \n", + "\n", + ] + ) + filename = f"{path}/{filename}" + with open(filename, "w") as f: + f.write(source_head) + f.write(source_body) + f.write(source_foot) + + +class NaDimDispatcher: + def __init__( + self, + ): + self.name = "DISPATCH_HOPPER_FNA_BACKWARD" + self.dims = [] + + def append(self, na_dim: int): + self.dims.append(na_dim) + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += f"#define {self.name}(rank, dtype, dim, is_causal, q_tile_shape, kv_tile_shape, ...) \\\n" + dispatcher_str += " [&] { \\\n" + for i, na_dim in enumerate(self.dims): + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + dispatcher_str += f"if constexpr (rank == {na_dim})" + dispatcher_str += " { \\\n" + dispatcher_str += " " + dispatcher_str += f" {self.name}_{na_dim}D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \\\n" + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += ( + ' throw std::runtime_error("Hopper FNA backward kernel dispatch failed! ' + + "It only supports NA1D, 2D, and 3D!" + + '"); \\\n' + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class DTypeDispatcher: + def __init__(self, na_dim: int): + self.dtypes: List[DataType] = [] + self.na_dim = na_dim + self.name = f"DISPATCH_HOPPER_FNA_BACKWARD_{self.na_dim}D" + + def append(self, dtype: DataType): + self.dtypes.append(dtype) + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += f"#define {self.name}(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, ...) \\\n" + dispatcher_str += " [&] { \\\n" + for i, dtype in enumerate(self.dtypes): + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + dispatcher_str += f"if (dtype == {dtype.torch_name})" + dispatcher_str += " { \\\n" + dispatcher_str += " " + dispatcher_str += f" {self.name}_{dtype.short_name}(dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \\\n" + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += ( + ' throw std::runtime_error("Hopper FNA-' + + f"{self.na_dim}D backward kernel dispatch failed! " + 'It does not support dtype "' + ' + std::string(c10::toString(dtype)) + "."' + "); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class HeadDimDispatcher: + def __init__(self, na_dim: int, dtype: DataType): + self.na_dim = na_dim + self.dtype = dtype + self.name = ( + f"DISPATCH_HOPPER_FNA_BACKWARD_{self.na_dim}D_{self.dtype.short_name}" + ) + + self.dims: List[int] = [] + + def append(self, dim: int): + self.dims.append(dim) + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += f"#define {self.name}(dim, is_causal, q_tile_shape, kv_tile_shape, ...) \\\n" + dispatcher_str += " [&] { \\\n" + for i, dim in enumerate(self.dims): + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + dispatcher_str += f"if (dim == {dim})" + dispatcher_str += " { \\\n" + dispatcher_str += " " + dispatcher_str += f" {self.name}_headdim{dim}(is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \\\n" + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += ( + " throw std::runtime_error(" + '"Hopper FNA backward kernel dispatch failed! ' + 'It does not support head dim "' + + f' + std::to_string(dim) + " for {self.dtype.short_name}."' + + "); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class CausalMaskDispatcher: + def __init__( + self, + na_dim: int, + dtype: DataType, + head_dim: int, + ): + self.na_dim = na_dim + self.dtype = dtype + self.head_dim = head_dim + + self.name = f"DISPATCH_HOPPER_FNA_BACKWARD_{self.na_dim}D_{self.dtype.short_name}_headdim{head_dim}" + + self.cms: List = [] + + def append(self, cm): + assert len(cm) == self.na_dim, f"{cm} incompatible for {self.na_dim}D NA." + self.cms.append(cm) + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += ( + f"#define {self.name}(is_causal, q_tile_shape, kv_tile_shape, ...) \\\n" + ) + dispatcher_str += " [&] { \\\n" + i = 0 + for cm in self.cms: + cm_str = "causal" + "x".join(["1" if c else "0" for c in cm]) + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + i += 1 + dispatcher_str += "if (" + for dim in range(self.na_dim): + dispatcher_str += ( + f"cute::get<{dim}>(is_causal)" + if cm[dim] + else f"not cute::get<{dim}>(is_causal)" + ) + if dim != self.na_dim - 1: + dispatcher_str += " && " + dispatcher_str += ")" + dispatcher_str += " { \\\n" + dispatcher_str += " " + dispatcher_str += f" {self.name}_{cm_str}(q_tile_shape, kv_tile_shape, __VA_ARGS__); \\\n" + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += " " + dispatcher_str += ( + ' throw std::runtime_error("Hopper FNA-' + + f"{self.na_dim}D backward kernel dispatch failed! " + f"Causal mask dispatcher ({self.dtype.short_name}, " + f'head_dim {self.head_dim}) got invalid causal mask!"' + "); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class ConfigDispatcher: + def __init__( + self, + na_dim: int, + dtype: DataType, + head_dim: int, + causal_mask: tuple, + ): + self.na_dim = na_dim + self.dtype = dtype + self.head_dim = head_dim + self.causal_mask = causal_mask + cm_str = "_causal" + "x".join(["1" if c else "0" for c in self.causal_mask]) + + self.name = f"DISPATCH_HOPPER_FNA_BACKWARD_{self.na_dim}D_{self.dtype.short_name}_headdim{head_dim}" + self.name += cm_str + + self.configs: List = [] + + def append(self, config): + assert len(config) == 2 + assert isinstance(config[0], tuple) and len(config[0]) == self.na_dim + assert isinstance(config[1], tuple) and len(config[1]) == self.na_dim + self.configs.append(config) + + def get_kernel_instance(self, q_tile_shape, kv_tile_shape): + gemm_M = math.prod(q_tile_shape) + gemm_N = math.prod(kv_tile_shape) + gemm_K = self.head_dim + gemm_shape = (gemm_M, gemm_N, gemm_K) + + supported_configs = SUPPORTED_CONFIGS_BACKWARD[self.dtype.bits][self.head_dim] + assert ( + gemm_shape in supported_configs + ), f"{gemm_shape=} not in supported configs {supported_configs=}" + + kernel = HopperFnaInstance( + na_dim=self.na_dim, + dtype=self.dtype, + gemm_shape=gemm_shape, + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + causal=self.causal_mask, + ) + + return kernel + + def get_target_name(self, q_tile_shape, kv_tile_shape): + kernel = self.get_kernel_instance(q_tile_shape, kv_tile_shape) + return kernel.get_name() + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += f"#define {self.name}(q_tile_shape, kv_tile_shape, ...) \\\n" + dispatcher_str += " [&] { \\\n" + i = 0 + for q_tile_shape, kv_tile_shape in self.configs: + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + i += 1 + dispatcher_str += "if (" + for dim in range(self.na_dim): + dispatcher_str += ( + f"cute::get<{dim}>(q_tile_shape) == {q_tile_shape[dim]}" + ) + dispatcher_str += " && " + dispatcher_str += "\\\n" + for dim in range(self.na_dim): + dispatcher_str += ( + f"cute::get<{dim}>(kv_tile_shape) == {kv_tile_shape[dim]}" + ) + if dim < self.na_dim - 1: + dispatcher_str += " && " + dispatcher_str += ")" + dispatcher_str += " { \\\n" + + dispatcher_str += f" natten::cuda::fna_hopper::{self.get_target_name(q_tile_shape, kv_tile_shape)}(__VA_ARGS__); \\\n" + + dispatcher_str += "} \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += " " + q_tile_str = ' + "," + '.join( + [ + f"std::to_string(cute::get<{i}>(q_tile_shape))" + for i in range(self.na_dim) + ] + ) + kv_tile_str = ' + "," + '.join( + [ + f"std::to_string(cute::get<{i}>(kv_tile_shape))" + for i in range(self.na_dim) + ] + ) + dispatcher_str += ( + ' throw std::runtime_error("Hopper FNA-' + + f"{self.na_dim}D backward kernel dispatch failed! " + "It got invalid Q tile and KV tile " + + f"combination ({self.dtype.short_name}, " + f"head_dim {self.head_dim}): " + 'q_tile=(" + ' + q_tile_str + ' + "), ' + 'kv_tile=(" + ' + kv_tile_str + ' + ")."' + "); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +def write_header_file(content, path, namespaces, extra_includes=None): + extra_includes = extra_includes or [] + header_head = [ + "#pragma once\n", + "\n\n", + ] + header_head += ["#include \n"] + header_head += ["#include \n"] + header_head += ["#include \n"] + header_head += ["#include \n"] + header_head += ["#ifdef NATTEN_WITH_CUTLASS\n"] + header_head += ["#ifdef NATTEN_WITH_HOPPER_FNA\n"] + for incl in extra_includes: + header_head += [f"#include <{incl}> \n"] + + for namespace in namespaces: + header_head += [f"namespace {namespace}", " { \n"] + + header_foot = [ + "\n\n", + ] + for namespace in namespaces: + header_foot += ["} ", f"// namespace {namespace}", " \n"] + header_foot += [ + "#endif \n", + "#endif \n", + "\n", + ] + with open(path, "w") as f: + f.write("".join(header_head)) + f.write(content) + f.write("".join(header_foot)) + + +def generate_hopper_fna_kernels(path, num_splits=2): + + NA_DIMS = [1, 2, 3] + + SUPPORTED_DTYPES = [ + Half, + BFloat, + ] + + HEAD_DIMS = [32, 64, 128] + + CAUSAL_MASKS = { + 1: [(False,), (True,)], + 2: [(False, False), (False, True), (True, False), (True, True)], + 3: [ + (False, False, False), + (False, False, True), + (False, True, False), + (False, True, True), + (True, False, False), + (True, False, True), + (True, True, False), + (True, True, True), + ], + } + + CONFIGS = { + 1: { + 16: { + 32: [ + ((64,), (128,)), + ((128,), (128,)), + ], + 64: [ + ((64,), (128,)), + ((128,), (128,)), + ], + 128: [ + ((64,), (128,)), + ], + }, + }, + 2: { + 16: { + 32: [ + ((8, 8), (16, 8)), + ((8, 8), (8, 16)), + ((16, 8), (16, 8)), + ((16, 8), (8, 16)), + ], + 64: [ + ((8, 8), (16, 8)), + ((8, 8), (8, 16)), + ((16, 8), (16, 8)), + ((16, 8), (8, 16)), + ], + 128: [ + ((8, 8), (16, 8)), + ((8, 8), (8, 16)), + ], + }, + }, + 3: { + 16: { + 32: [ + ((4, 4, 4), (4, 4, 8)), + ((4, 4, 4), (2, 8, 8)), + ((4, 4, 8), (4, 4, 8)), + ((4, 4, 8), (2, 8, 8)), + ], + 64: [ + ((4, 4, 4), (4, 4, 8)), + ((4, 4, 4), (2, 8, 8)), + ((4, 4, 8), (4, 4, 8)), + ((4, 4, 8), (2, 8, 8)), + ], + 128: [ + ((4, 4, 4), (4, 4, 8)), + ((4, 4, 4), (2, 8, 8)), + ((2, 4, 8), (2, 8, 8)), + ((1, 8, 8), (2, 8, 8)), + ], + }, + }, + } + + dtype_dispatchers = [] + head_dim_dispatchers = [] + cm_dispatchers = [] + config_dispatchers = [] + kernels = [] + + rank_dispatcher = NaDimDispatcher() + for na_dim in NA_DIMS: + rank_dispatcher.append(na_dim) + + dtype_dispatcher = DTypeDispatcher(na_dim=na_dim) + for dtype in SUPPORTED_DTYPES: + dtype_dispatcher.append(dtype) + + head_dim_dispatcher = HeadDimDispatcher(dtype=dtype, na_dim=na_dim) + + for head_dim in HEAD_DIMS: + head_dim_dispatcher.append(head_dim) + + cm_dispatcher = CausalMaskDispatcher( + na_dim=na_dim, dtype=dtype, head_dim=head_dim + ) + for cm in CAUSAL_MASKS[na_dim]: + cm_dispatcher.append(cm) + + config_dispatcher = ConfigDispatcher( + dtype=dtype, + head_dim=head_dim, + na_dim=na_dim, + causal_mask=cm, + ) + + for q_tile_shape, kv_tile_shape in CONFIGS[na_dim][dtype.bits][ + head_dim + ]: + config_dispatcher.append((q_tile_shape, kv_tile_shape)) + kernels.append( + config_dispatcher.get_kernel_instance( + q_tile_shape, kv_tile_shape + ) + ) + config_dispatchers.append(config_dispatcher) + cm_dispatchers.append(cm_dispatcher) + head_dim_dispatchers.append(head_dim_dispatcher) + dtype_dispatchers.append(dtype_dispatcher) + + # + + path_to_sources = f"{path}/autogen/src/cuda/hopper_fna_bwd/" + rel_header = "natten_autogen/cuda/hopper_fna_bwd/" + path_to_header_dir = f"{path}/autogen/include/{rel_header}" + + os.makedirs(path_to_sources, exist_ok=False) + os.makedirs(path_to_header_dir, exist_ok=False) + + path_headers = f"{path_to_header_dir}kernels.h" + path_rank = f"{path_to_header_dir}interface.h" + path_dtype = f"{path_to_header_dir}dispatch_dtype.h" + path_head_dim = f"{path_to_header_dir}dispatch_head_dim.h" + path_cm = f"{path_to_header_dir}dispatch_cm.h" + path_tile_shape = f"{path_to_header_dir}dispatch_tile_shape.h" + + rel_path_headers = f"{rel_header}kernels.h" + rel_path_dtype = f"{rel_header}dispatch_dtype.h" + rel_path_head_dim = f"{rel_header}dispatch_head_dim.h" + rel_path_cm = f"{rel_header}dispatch_cm.h" + rel_path_tile_shape = f"{rel_header}dispatch_tile_shape.h" + + rank_disp = rank_dispatcher.get_dispatcher() + + dtype_disp = "" + for dispatcher in dtype_dispatchers: + dtype_disp += dispatcher.get_dispatcher() + + head_dim_disp = "" + for dispatcher in head_dim_dispatchers: + head_dim_disp += dispatcher.get_dispatcher() + + cm_disp = "" + for dispatcher in cm_dispatchers: + cm_disp += dispatcher.get_dispatcher() + + config_disp = "" + for dispatcher in config_dispatchers: + config_disp += dispatcher.get_dispatcher() + + headers = "" + for kernel in kernels: + headers += kernel.get_decl() + + assert ( + len(kernels) >= num_splits + ), f"Generated {len(kernels)} kernels, but got {num_splits=}." + split_size = len(kernels) // num_splits + num_splits_with_res = len(kernels) % num_splits + kernels_emitted = [] + kernels_split = [] + for split_idx in range(num_splits): + kernel_start_idx = split_size * split_idx + min(num_splits_with_res, split_idx) + num_kernels_in_split = split_size + ( + 1 if split_idx < num_splits_with_res else 0 + ) + kernel_end_idx = kernel_start_idx + num_kernels_in_split + assert kernel_end_idx <= len(kernels) + pth_set = set() + source_list = [] + for kernel_idx in range(kernel_start_idx, kernel_end_idx): + kernel = kernels[kernel_idx] + # pth_set.add(kernel.path_to_header) + source_list.append(kernel) + kernels_emitted.append(kernel_idx) + pth_set.add(rel_path_headers) + write_combined_source_file( + path_to_sources, f"source_{split_idx}.cu", sorted(pth_set), source_list + ) + kernels_split.append(source_list) + # print(f"{split_idx=}, {kernel_start_idx=}, {kernel_end_idx=}, {len(kernels_emitted)=}") + assert split_idx == num_splits - 1, f"Expected {split_idx=} == {num_splits=} - 1" + assert len(kernels_emitted) == len(kernels) + assert sorted(kernels_emitted) == [ + x for x in range(len(kernels)) + ], f"{sorted(kernels_emitted)=}" + assert all(len(x) > 0 for x in kernels_split) + + namespaces = ["natten", "cuda", "fna_hopper"] + cuda_headers = [ + "natten/natten.h", + "ATen/ATen.h", + "ATen/cuda/CUDAContext.h", + "c10/cuda/CUDAGuard.h", + "c10/cuda/CUDAStream.h", + "natten/natten.h", + "natten/helpers.h", + "natten/cuda/fna_hopper/fna_backward.cuh", + ] + write_header_file(rank_disp, path_rank, namespaces, cuda_headers + [rel_path_dtype]) + write_header_file( + dtype_disp, path_dtype, namespaces, cuda_headers + [rel_path_head_dim] + ) + write_header_file( + head_dim_disp, path_head_dim, namespaces, cuda_headers + [rel_path_cm] + ) + write_header_file( + cm_disp, path_cm, namespaces, cuda_headers + [rel_path_tile_shape] + ) + write_header_file( + config_disp, path_tile_shape, namespaces, cuda_headers + [rel_path_headers] + ) + write_header_file(headers, path_headers, namespaces, cuda_headers) + + +def generate_hopper_fna(output_directory: str, num_splits: int): + generate_hopper_fna_kernels(output_directory, num_splits=num_splits) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "-o", + "--output-directory", + default=DEFAULT_OUTPUT_DIR, + help="Path to the directory where the auto-generated " + "kernel instantiations are dumped. " + f"Default: {DEFAULT_OUTPUT_DIR}", + ) + parser.add_argument( + "--num-splits", + type=int, + default=8, + help="Number of source files into which the kernels are split. Default: 8.", + ) + args = parser.parse_args() + generate_hopper_fna(args.output_directory, args.num_splits) diff --git a/natten/scripts/autogen_reference_fna.py b/natten/scripts/autogen_reference_fna.py new file mode 100644 index 00000000..077bfd73 --- /dev/null +++ b/natten/scripts/autogen_reference_fna.py @@ -0,0 +1,621 @@ +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# This script is intended to emit fused kernel instantiations into +# a variable number of source files, generate appropriate headers +# and a single dispatcher interface, which will be used by the +# NATTEN API to call the kernels. +# +# NOTE: these scripts are heavily under-documented, and +# overly-repetitive, and will be replaced in future PRs. +# Please use it with caution. + +import argparse +import os +from typing import List + +DEFAULT_OUTPUT_DIR = "csrc/" + + +KERNEL_DECL_TEMPLATE = """ +void {kernel_name}( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + {DimType} qkv_shape, + {DimType} window_size, + {DimType} stride, + {DimType} dilation, + float attn_scale, + cudaStream_t stream); +""" + + +KERNEL_IMPL_TEMPLATE = """ +void {kernel_name}( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + {DimType} qkv_shape, + {DimType} window_size, + {DimType} stride, + {DimType} dilation, + float attn_scale, + cudaStream_t stream) {{ + + using Causal = {Causal}; + + fna_reference_forward( + static_cast<{dtype}*>(ptr_Q), + static_cast<{dtype}*>(ptr_K), + static_cast<{dtype}*>(ptr_V), + static_cast<{dtype}*>(ptr_O), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{{}}, + attn_scale, + stream); +}} +""" + + +KERNEL_BWD_DECL_TEMPLATE = """ +void {kernel_name}( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + {DimType} qkv_shape, + {DimType} window_size, + {DimType} stride, + {DimType} dilation, + float attn_scale, + cudaStream_t stream); +""" + + +KERNEL_BWD_IMPL_TEMPLATE = """ +void {kernel_name}( + void* ptr_Q, + void* ptr_K, + void* ptr_V, + void* ptr_O, + void* ptr_DO, + void* ptr_DQ, + void* ptr_DK, + void* ptr_DV, + void* ptr_LSE, + int batch_size, + int seqlen, + int heads, + int heads_kv, + int dim, + int dim_value, + int num_additional_kv, + {DimType} qkv_shape, + {DimType} window_size, + {DimType} stride, + {DimType} dilation, + float attn_scale, + cudaStream_t stream) {{ + + using Causal = {Causal}; + + fna_reference_backward( + static_cast<{dtype}*>(ptr_Q), + static_cast<{dtype}*>(ptr_K), + static_cast<{dtype}*>(ptr_V), + static_cast<{dtype}*>(ptr_O), + static_cast<{dtype}*>(ptr_DO), + static_cast<{dtype}*>(ptr_DQ), + static_cast<{dtype}*>(ptr_DK), + static_cast<{dtype}*>(ptr_DV), + static_cast(ptr_LSE), + batch_size, + seqlen, + heads, + heads_kv, + dim, + dim_value, + num_additional_kv, + qkv_shape, + window_size, + stride, + dilation, + Causal{{}}, + attn_scale, + stream); +}} +""" + + +class DataType: + def __init__(self, name, short_name, torch_name, bits): + self.name = name + self.bits = bits + self.short_name = short_name + self.torch_name = torch_name + + +Float = DataType("float", "float32", "torch::kFloat32", 32) +Half = DataType("cutlass::half_t", "float16", "torch::kFloat16", 16) +BFloat = DataType("cutlass::bfloat16_t", "bfloat16", "torch::kBFloat16", 16) + + +def iterable_to_static_cute_tuple(shape_in) -> str: + shape = ", ".join([f"cute::Int<{x}>" for x in shape_in]) + return f"cute::tuple<{shape}>" + + +def get_dim_type(na_dim: int) -> str: + shape = ", ".join(["int" for _ in range(na_dim)]) + return f"cute::tuple<{shape}>" + + +class ReferenceFnaInstance: + def __init__( + self, + na_dim: int, + dtype: DataType, + causal: tuple, + is_backward: bool, + ): + assert 0 < na_dim <= 3 + assert na_dim == len(causal) + self.na_dim = na_dim + self.causal = causal + self.dtype = dtype + self.is_backward = is_backward + + def get_causal_cute(self) -> str: + consts = ", ".join( + ["cute::true_type" if c else "cute::false_type" for c in self.causal] + ) + return f"cute::tuple<{consts}>" + + def get_name(self) -> str: + backward_str = "" if not self.is_backward else "_backward" + name = f"reference_fna{self.na_dim}d{backward_str}" + name += f"_{self.dtype.short_name}" + name += "_causal" + "x".join(["1" if c else "0" for c in self.causal]) + return name + + def get_decl(self) -> str: + return ( + KERNEL_BWD_DECL_TEMPLATE if self.is_backward else KERNEL_DECL_TEMPLATE + ).format( + kernel_name=self.get_name(), + DimType=get_dim_type(self.na_dim), + dtype=self.dtype.name, + ) + + def get_impl(self) -> str: + return ( + KERNEL_BWD_IMPL_TEMPLATE if self.is_backward else KERNEL_IMPL_TEMPLATE + ).format( + kernel_name=self.get_name(), + DimType=get_dim_type(self.na_dim), + Causal=self.get_causal_cute(), + dtype=self.dtype.name, + ) + + +def write_combined_source_file(path, filename, headers, kernels): + source_head = [] + source_head += ["#ifdef NATTEN_WITH_CUTLASS\n"] + + source_head += ["#include \n"] + source_head += ["#include \n"] + + source_head += ["#include \n"] + source_head += ["#include \n"] + source_head += ["#include \n"] + source_head += ["#include \n"] + + source_head += ["#include \n"] + source_head += ["#include \n"] + + source_head += ["#include \n"] + source_head += ["#include \n"] + + for header in headers: + source_head += [f"#include <{header}>\n"] + + source_head += ["namespace natten { \n"] + source_head += ["namespace cuda { \n"] + source_head += ["namespace reference { \n\n"] + + source_head = "".join(source_head) + + source_body = [] + for kernel in kernels: + source_body += "\n\n" + kernel.get_impl() + "\n\n" + source_body = "".join(source_body) + + source_foot = "".join( + [ + "} // namespace reference \n", + "} // namespace cuda \n", + "} // namespace natten \n", + "#endif \n", + "\n", + ] + ) + filename = f"{path}/{filename}" + with open(filename, "w") as f: + f.write(source_head) + f.write(source_body) + f.write(source_foot) + + +class NaDimDispatcher: + def __init__(self, is_backward: bool): + fwd_bwd_str = "BACKWARD" if is_backward else "FORWARD" + self.name = f"DISPATCH_REFERENCE_FNA_{fwd_bwd_str}" + self.dims: List[int] = [] + self.is_backward = is_backward + + def append(self, na_dim: int): + self.dims.append(na_dim) + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += f"#define {self.name}(rank, dtype, is_causal, ...) \\\n" + dispatcher_str += " [&] { \\\n" + for i, na_dim in enumerate(self.dims): + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + dispatcher_str += f"if constexpr (rank == {na_dim})" + dispatcher_str += " { \\\n" + dispatcher_str += " " + dispatcher_str += ( + f" {self.name}_{na_dim}D(dtype, is_causal, __VA_ARGS__); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += ( + ' throw std::runtime_error("Reference FNA kernel dispatch failed! ' + + "NATTEN only supports NA1D, 2D, and 3D!" + + '"); \\\n' + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class DTypeDispatcher: + def __init__(self, is_backward: bool, na_dim: int): + self.dtypes: List[DataType] = [] + self.na_dim = na_dim + fwd_bwd_str = "BACKWARD" if is_backward else "FORWARD" + self.name = f"DISPATCH_REFERENCE_FNA_{fwd_bwd_str}_{self.na_dim}D" + self.is_backward = is_backward + + def append(self, dtype: DataType): + self.dtypes.append(dtype) + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += f"#define {self.name}(dtype, is_causal, ...) \\\n" + dispatcher_str += " [&] { \\\n" + for i, dtype in enumerate(self.dtypes): + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + dispatcher_str += f"if (dtype == {dtype.torch_name})" + dispatcher_str += " { \\\n" + dispatcher_str += " " + dispatcher_str += ( + f" {self.name}_{dtype.short_name}(is_causal, __VA_ARGS__); \\\n" + ) + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += ( + ' throw std::runtime_error("Reference FNA kernel dispatch failed! ' + + f"Reference FNA-{self.na_dim}D does not support this data type." + + '"); \\\n' + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +class CausalMaskDispatcher: + def __init__( + self, + is_backward: bool, + na_dim: int, + dtype: DataType, + ): + self.na_dim = na_dim + self.dtype = dtype + + fwd_bwd_str = "BACKWARD" if is_backward else "FORWARD" + self.name = f"DISPATCH_REFERENCE_FNA_{fwd_bwd_str}_{self.na_dim}D_{self.dtype.short_name}" + self.is_backward = is_backward + + self.cms: List = [] + + def append(self, cm): + assert len(cm) == self.na_dim, f"{cm} incompatible for {self.na_dim}D NA." + self.cms.append(cm) + + def get_kernel_instance(self, cm): + return ReferenceFnaInstance( + na_dim=self.na_dim, + dtype=self.dtype, + causal=cm, + is_backward=self.is_backward, + ) + + def get_target_name(self, cm): + kernel = self.get_kernel_instance(cm) + return kernel.get_name() + + def get_dispatcher(self): + dispatcher_str = "" + dispatcher_str += f"#define {self.name}(is_causal, ...) \\\n" + dispatcher_str += " [&] { \\\n" + i = 0 + for cm in self.cms: + dispatcher_str += " " + if i > 0: + dispatcher_str += "else " + i += 1 + dispatcher_str += "if (" + for dim in range(self.na_dim): + dispatcher_str += ( + f"cute::get<{dim}>(is_causal)" + if cm[dim] + else f"not cute::get<{dim}>(is_causal)" + ) + if dim != self.na_dim - 1: + dispatcher_str += " && " + dispatcher_str += ")" + dispatcher_str += " { \\\n" + dispatcher_str += " " + + dispatcher_str += f" natten::cuda::reference::{self.get_target_name(cm)}(__VA_ARGS__); \\\n" + + dispatcher_str += " } \\\n" + dispatcher_str += " else { \\\n" + dispatcher_str += " " + dispatcher_str += ( + ' throw std::runtime_error("Reference FNA kernel dispatch failed! ' + + "Causal mask dispatcher got invalid causal mask!" + + '"); \\\n' + ) + dispatcher_str += " } \\\n" + dispatcher_str += "}();" + dispatcher_str += "\n\n" + return dispatcher_str + + +def write_header_file(content, path, namespaces, extra_includes=None): + extra_includes = extra_includes or [] + header_head = [ + "#pragma once\n", + "\n\n", + ] + header_head += ["#include \n"] + header_head += ["#include \n"] + header_head += ["#include \n"] + header_head += ["#ifdef NATTEN_WITH_CUTLASS\n"] + for incl in extra_includes: + header_head += [f"#include <{incl}> \n"] + + for namespace in namespaces: + header_head += [f"namespace {namespace}", " { \n"] + + header_foot = [ + "\n\n", + ] + for namespace in namespaces: + header_foot += ["} ", f"// namespace {namespace}", " \n"] + header_foot += [ + "#endif \n", + "\n", + ] + with open(path, "w") as f: + f.write("".join(header_head)) + f.write(content) + f.write("".join(header_foot)) + + +def generate_reference_fna_kernels(path, num_splits=2): + + NA_DIMS = [1, 2, 3] + + SUPPORTED_DTYPES = [ + Float, + Half, + BFloat, + ] + + CAUSAL_MASKS = { + 1: [(False,), (True,)], + 2: [(False, False), (False, True), (True, False), (True, True)], + 3: [ + (False, False, False), + (False, False, True), + (False, True, False), + (False, True, True), + (True, False, False), + (True, False, True), + (True, True, False), + (True, True, True), + ], + } + + dtype_dispatchers = [] + cm_dispatchers = [] + kernels = [] + + rank_dispatcher = NaDimDispatcher(is_backward=False) + rank_dispatcher_bwd = NaDimDispatcher(is_backward=True) + for na_dim in NA_DIMS: + rank_dispatcher.append(na_dim) + rank_dispatcher_bwd.append(na_dim) + + for is_backward in [False, True]: + dtype_dispatcher = DTypeDispatcher(is_backward=is_backward, na_dim=na_dim) + for dtype in SUPPORTED_DTYPES: + dtype_dispatcher.append(dtype) + + cm_dispatcher = CausalMaskDispatcher( + is_backward=is_backward, na_dim=na_dim, dtype=dtype + ) + for cm in CAUSAL_MASKS[na_dim]: + cm_dispatcher.append(cm) + kernels.append(cm_dispatcher.get_kernel_instance(cm)) + + cm_dispatchers.append(cm_dispatcher) + dtype_dispatchers.append(dtype_dispatcher) + + # + + path_to_sources = f"{path}/autogen/src/cuda/reference/" + rel_header = "natten_autogen/cuda/reference/" + path_to_header_dir = f"{path}/autogen/include/{rel_header}" + + os.makedirs(path_to_sources, exist_ok=False) + os.makedirs(path_to_header_dir, exist_ok=False) + + path_headers = f"{path_to_header_dir}kernels.h" + path_rank = f"{path_to_header_dir}interface.h" + path_dtype = f"{path_to_header_dir}dispatch_dtype.h" + path_cm = f"{path_to_header_dir}dispatch_cm.h" + + rel_path_headers = f"{rel_header}kernels.h" + rel_path_dtype = f"{rel_header}dispatch_dtype.h" + rel_path_cm = f"{rel_header}dispatch_cm.h" + + rank_disp = rank_dispatcher.get_dispatcher() + rank_dispatcher_bwd.get_dispatcher() + + dtype_disp = "" + for dispatcher in dtype_dispatchers: + dtype_disp += dispatcher.get_dispatcher() + + cm_disp = "" + for dispatcher in cm_dispatchers: + cm_disp += dispatcher.get_dispatcher() + + headers = "" + for kernel in kernels: + headers += kernel.get_decl() + + assert ( + len(kernels) >= num_splits + ), f"Generated {len(kernels)} kernels, but got {num_splits=}." + split_size = len(kernels) // num_splits + num_splits_with_res = len(kernels) % num_splits + kernels_emitted = [] + kernels_split = [] + for split_idx in range(num_splits): + kernel_start_idx = split_size * split_idx + min(num_splits_with_res, split_idx) + num_kernels_in_split = split_size + ( + 1 if split_idx < num_splits_with_res else 0 + ) + kernel_end_idx = kernel_start_idx + num_kernels_in_split + assert kernel_end_idx <= len(kernels) + pth_set = set() + source_list = [] + for kernel_idx in range(kernel_start_idx, kernel_end_idx): + kernel = kernels[kernel_idx] + source_list.append(kernel) + kernels_emitted.append(kernel_idx) + pth_set.add(rel_path_headers) + write_combined_source_file( + path_to_sources, f"source_{split_idx}.cu", sorted(pth_set), source_list + ) + kernels_split.append(source_list) + # print(f"{split_idx=}, {kernel_start_idx=}, {kernel_end_idx=}, {len(kernels_emitted)=}") + assert split_idx == num_splits - 1, f"Expected {split_idx=} == {num_splits=} - 1" + assert len(kernels_emitted) == len(kernels) + assert sorted(kernels_emitted) == [ + x for x in range(len(kernels)) + ], f"{sorted(kernels_emitted)=}" + assert all(len(x) > 0 for x in kernels_split) + + namespaces = ["natten", "cuda", "reference"] + cuda_headers = [ + "natten/natten.h", + "ATen/ATen.h", + "ATen/cuda/CUDAContext.h", + "c10/cuda/CUDAGuard.h", + "c10/cuda/CUDAStream.h", + "natten/natten.h", + "natten/helpers.h", + "natten/cuda/reference/fna_reference_forward.hpp", + "natten/cuda/reference/fna_reference_backward.hpp", + ] + write_header_file(rank_disp, path_rank, namespaces, cuda_headers + [rel_path_dtype]) + write_header_file(dtype_disp, path_dtype, namespaces, cuda_headers + [rel_path_cm]) + write_header_file(cm_disp, path_cm, namespaces, cuda_headers + [rel_path_headers]) + write_header_file(headers, path_headers, namespaces, cuda_headers) + + +def generate_reference_fna(output_directory: str, num_splits: int): + generate_reference_fna_kernels(output_directory, num_splits=num_splits) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "-o", + "--output-directory", + default=DEFAULT_OUTPUT_DIR, + help="Path to the directory where the auto-generated " + "kernel instantiations are dumped. " + f"Default: {DEFAULT_OUTPUT_DIR}", + ) + parser.add_argument( + "--num-splits", + type=int, + default=2, + help="Number of source files into which the kernels are split. Default: 2.", + ) + args = parser.parse_args() + generate_reference_fna(args.output_directory, args.num_splits) diff --git a/natten/scripts/generate_build_toml.py b/natten/scripts/generate_build_toml.py new file mode 100644 index 00000000..799f0493 --- /dev/null +++ b/natten/scripts/generate_build_toml.py @@ -0,0 +1,145 @@ +# Regenerates build.toml's kernel sections from the files on disk. +# +# kernel-builder requires static `src` lists (they define both the Nix source +# fileset and the CMake sources), while NATTEN's kernel instantiations are +# autogenerated. Run this after `scripts/regen.sh` so the lists stay in sync +# with the autogen output. Usage: +# +# python3 scripts/generate_build_toml.py + +import os +import sys +from pathlib import Path + +KERNEL_DIR = Path(__file__).resolve().parent.parent + +HEADER = """\ +[general] +name = "natten" +version = 1 +edition = 5 +license = "MIT" +upstream = "https://github.com/SHI-Labs/NATTEN" +backends = ["cuda"] + +[general.hub] +repo-id = "kernels-community/natten" + +[torch] +src = ["torch-ext/torch_binding.cpp"] +include = ["csrc/include"] +""" + +COMMON_CUDA_FLAGS = [ + "-O3", + "-std=c++17", + "--expt-relaxed-constexpr", + "--extended-lambda", + "--use_fast_math", + "--ftemplate-backtrace-limit=0", + "-Xfatbin=-compress-all", + "-DNDEBUG", + "-DNATTEN_WITH_CUTLASS", + "-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1", +] + +DEPENDS = ["torch", "cutlass_4_5"] +INCLUDE = ["csrc/include", "csrc/autogen/include"] + +# Upstream compiles hopper (90) and blackwell (100/103) kernels with +# arch-specific ("a") targets; see setup.py's arch_list_to_cmake_tags. +# SM103 is not in kernel-builder's supported arch list yet, so Blackwell +# Ultra (B300) is not covered. Blackwell kernels need CUDA >= 12.8; the +# section is skipped (via cuda-minver) on older toolkits, and the natten.h +# CUDART_VERSION gate turns the static dispatchers into stubs to match. +SECTIONS = [ + { + "name": "natten", + "capabilities": ["7.0", "7.5", "8.0", "8.6", "8.9", "9.0", "10.0", "12.0"], + "minver": None, + "extra_flags": [], + "autogen_dirs": ["reference", "fna", "fmha"], + "static_src": True, + }, + { + "name": "natten_hopper", + "capabilities": ["9.0a"], + "minver": None, + "extra_flags": ["-DNATTEN_WITH_HOPPER_FNA=1"], + "autogen_dirs": [ + "hopper_fna", + "hopper_fna_bwd", + "hopper_fmha", + "hopper_fmha_bwd", + ], + "static_src": False, + }, + { + "name": "natten_blackwell", + "capabilities": ["10.0a"], + "minver": "12.8", + "extra_flags": ["-DNATTEN_WITH_BLACKWELL_FNA=1"], + "autogen_dirs": [ + "blackwell_fna", + "blackwell_fna_bwd", + "blackwell_fmha", + "blackwell_fmha_bwd", + ], + "static_src": False, + }, +] + + +def rel_files(subdir: str, exts=(".h", ".hpp", ".cuh", ".cu", ".cpp")) -> list: + root = KERNEL_DIR / subdir + if not root.is_dir(): + sys.exit(f"error: {root} does not exist; run scripts/regen.sh first") + out = [] + for path in sorted(root.rglob("*")): + if path.is_file() and path.suffix in exts: + out.append(str(path.relative_to(KERNEL_DIR))) + return out + + +def toml_list(key: str, values: list, indent: str = "") -> str: + lines = [f"{indent}{key} = ["] + for v in values: + lines.append(f'{indent} "{v}",') + lines.append(f"{indent}]") + return "\n".join(lines) + + +def main() -> None: + parts = [HEADER] + + for section in SECTIONS: + src = [] + if section["static_src"]: + # All shared headers ride along with the core section; every + # section adds csrc/include and csrc/autogen/include to its + # include path. + src += rel_files("csrc/include") + src += rel_files("csrc/src") + for autogen_dir in section["autogen_dirs"]: + src += rel_files(f"csrc/autogen/include/natten_autogen/cuda/{autogen_dir}") + src += rel_files(f"csrc/autogen/src/cuda/{autogen_dir}") + + lines = [f"[kernel.{section['name']}]"] + lines.append('backend = "cuda"') + if section["minver"]: + lines.append(f'cuda-minver = "{section["minver"]}"') + lines.append(toml_list("cuda-capabilities", section["capabilities"])) + lines.append( + toml_list("cuda-flags", COMMON_CUDA_FLAGS + section["extra_flags"]) + ) + lines.append(toml_list("depends", DEPENDS)) + lines.append(toml_list("include", INCLUDE)) + lines.append(toml_list("src", src)) + parts.append("\n".join(lines) + "\n") + + (KERNEL_DIR / "build.toml").write_text("\n".join(parts)) + print(f"wrote {KERNEL_DIR / 'build.toml'}") + + +if __name__ == "__main__": + main() diff --git a/natten/scripts/regen.sh b/natten/scripts/regen.sh new file mode 100755 index 00000000..9af47dbd --- /dev/null +++ b/natten/scripts/regen.sh @@ -0,0 +1,38 @@ +#!/usr/bin/env bash +# Regenerates the committed NATTEN autogen kernel instantiations and the +# build.toml src lists. Run this when bumping the vendored NATTEN version +# (after re-vendoring csrc/ and scripts/autogen_*.py from upstream). +# +# The split counts below are upstream's "default" autogen policy +# (NUM_SPLITS["default"] in NATTEN's setup.py). Note that the vendored +# autogen scripts are patched to not emit `#include ` +# (pybind11 is incompatible with the Python limited API that kernel-builder +# targets); re-apply that patch when re-vendoring the scripts. + +set -euo pipefail + +cd "$(dirname "$0")/.." + +rm -rf csrc/autogen + +python3 scripts/autogen_reference_fna.py --num-splits 2 -o csrc +python3 scripts/autogen_fna.py --num-splits 64 -o csrc +python3 scripts/autogen_fmha.py --num-splits 6 -o csrc +python3 scripts/autogen_hopper_fna.py --num-splits 8 -o csrc +python3 scripts/autogen_hopper_fna_bwd.py --num-splits 4 -o csrc +python3 scripts/autogen_hopper_fmha.py --num-splits 5 -o csrc +python3 scripts/autogen_hopper_fmha_bwd.py --num-splits 5 -o csrc +python3 scripts/autogen_blackwell_fna.py --num-splits 28 -o csrc +python3 scripts/autogen_blackwell_fna_bwd.py --num-splits 14 -o csrc +python3 scripts/autogen_blackwell_fmha.py --num-splits 4 -o csrc +python3 scripts/autogen_blackwell_fmha_bwd.py --num-splits 4 -o csrc + +if grep -rl "torch/extension.h" csrc/autogen >/dev/null; then + echo "error: autogen output includes torch/extension.h (pybind11);" >&2 + echo "the vendored autogen scripts lost the limited-API patch." >&2 + exit 1 +fi + +python3 scripts/generate_build_toml.py + +echo "Done. Review changes to csrc/autogen and build.toml, then commit." diff --git a/natten/tests/__init__.py b/natten/tests/__init__.py new file mode 100644 index 00000000..99f6c1e8 --- /dev/null +++ b/natten/tests/__init__.py @@ -0,0 +1,22 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# diff --git a/natten/tests/conftest.py b/natten/tests/conftest.py new file mode 100644 index 00000000..8397033e --- /dev/null +++ b/natten/tests/conftest.py @@ -0,0 +1,10 @@ +import pytest + +from .utils import logger + + +@pytest.fixture(autouse=True) +def log_test_name(request): + logger.debug(f"Starting {request.node.name}") + yield + logger.debug(f"Finished {request.node.name}") diff --git a/natten/tests/test_attn_merge.py b/natten/tests/test_attn_merge.py new file mode 100644 index 00000000..38b0a55e --- /dev/null +++ b/natten/tests/test_attn_merge.py @@ -0,0 +1,286 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +import random +import unittest +from typing import Tuple + +import natten # noqa: F401 +import torch +from natten._environment import _NUM_RAND_SWEEP_TESTS as RAND_SWEEP_TESTS +from natten.functional import attention, merge_attentions +from natten.utils.testing import ( + skip_if_blackwell_kernels_not_supported, + skip_if_hopper_kernels_not_supported, + skip_if_libnatten_is_not_supported, + skip_if_not_running_extended_tests, +) +from torch import Tensor + +from .utils import logger, reset_torch_compile + + +def _reset_everything(random_seed: int = 42, torch_seed: int = 42): + random.seed(random_seed) + torch.manual_seed(torch_seed) + logger.debug(f"Reset seeds: {random_seed=}, {torch_seed=}") + torch.cuda.empty_cache() + torch.use_deterministic_algorithms(False) + + +def sdpa_split( + q: Tensor, + k_list: list, + v_list: list, + do: Tensor, + backend: str, + torch_compile: bool, +) -> Tuple[Tensor, Tensor, Tensor, Tensor]: + q = q.requires_grad_(True) + k_list = [k.requires_grad_(True) for k in k_list] + v_list = [v.requires_grad_(True) for v in v_list] + + outputs, lses = [], [] + for k, v in zip(k_list, v_list): + out, lse = attention(q, k, v, return_lse=True, backend=backend) + outputs.append(out) + lses.append(lse) + + out, _ = merge_attentions(outputs, lses, torch_compile=torch_compile) + + out.backward(do) + + with torch.no_grad(): + output = out.data + assert q.grad is not None + dq = q.grad.data + + dk_list = [] + dv_list = [] + for k, v in zip(k_list, v_list): + assert k.grad is not None + assert v.grad is not None + dk_list.append(k.grad.data) + dv_list.append(v.grad.data) + + dk = torch.cat(dk_list, dim=1) + dv = torch.cat(dv_list, dim=1) + + return output, dq, dk, dv + + +def sdpa_ref( + q: Tensor, + k_list: list, + v_list: list, + do: Tensor, + backend: str, +) -> Tuple[Tensor, Tensor, Tensor, Tensor]: + with torch.no_grad(): + k = torch.cat(k_list, dim=1) + v = torch.cat(v_list, dim=1) + + q = q.requires_grad_(True) + k = k.requires_grad_(True) + v = v.requires_grad_(True) + + out: Tensor = attention(q, k, v, backend=backend) # type: ignore[assignment] + out.backward(do) + + assert q.grad is not None + assert k.grad is not None + assert v.grad is not None + + return out.data, q.grad.data, k.grad.data, v.grad.data + + +class AttentionMergeTest(unittest.TestCase): + def setUp(self): + _reset_everything() + + def tearDown(self): + _reset_everything() + + def _test( + self, + batch: int, + heads: int, + head_dim: int, + seqlen_Q: int, + seqlen_KV_list: list, + backend: str, + torch_compile: bool, + ): + + ALLOWED_DTYPES = [ + # (dtype, atol_out, (atol_dq, atol_dk, atol_dv)) + (torch.float32, 1e-3, (1e-2, 1e-3, 1e-3)), + (torch.float16, 1e-2, (1e-2, 1e-2, 1e-2)), + (torch.bfloat16, 5e-2, (5e-2, 5e-2, 5e-2)), + ] + + SUPPORTED_DTYPES = [torch.float16, torch.bfloat16] + # TODO: query backends directly for supported dtypes + if backend == "cutlass-fmha": + SUPPORTED_DTYPES += [torch.float32] + + for dtype, atol_out, (atol_dq, atol_dk, atol_dv) in ALLOWED_DTYPES: + if dtype not in SUPPORTED_DTYPES: + continue + + if torch_compile: + reset_torch_compile(1) + + logger.debug( + f"Testing Attention Merging: {batch=}, {heads=}, {head_dim=}, " + f"{seqlen_Q=}, seqlen_KV_list={seqlen_KV_list}, " + f"{dtype=}, {backend=}, {torch_compile=}." + ) + + q = torch.randn( + batch, seqlen_Q, heads, head_dim, device="cuda", dtype=dtype + ) + + k_list = [] + v_list = [] + for seqlen_KV in seqlen_KV_list: + k = torch.randn( + batch, seqlen_KV, heads, head_dim, device="cuda", dtype=dtype + ) + v = torch.randn_like(k) + k_list.append(k) + v_list.append(v) + + do = torch.randn_like(q) + + q_ref = q.clone() + k_list_ref = [k.clone() for k in k_list] + v_list_ref = [v.clone() for v in v_list] + do_ref = do.clone() + + output_ref, dq_ref, dk_ref, dv_ref = sdpa_ref( + q_ref, + k_list_ref, + v_list_ref, + do_ref, + backend=backend, + ) + + output, dq, dk, dv = sdpa_split( + q, k_list, v_list, do, backend=backend, torch_compile=torch_compile + ) + + torch.testing.assert_close( + output.float(), output_ref.float(), atol=atol_out, rtol=0 + ) + torch.testing.assert_close(dk.float(), dk_ref.float(), atol=atol_dk, rtol=0) + torch.testing.assert_close(dv.float(), dv_ref.float(), atol=atol_dv, rtol=0) + torch.testing.assert_close(dq.float(), dq_ref.float(), atol=atol_dq, rtol=0) + + def _test_randsweep(self, backend, num_tests=1000): + max_Q_ = 16384 + max_KV_total_ = 2**15 + for i in range(num_tests): + # to help with reproducibility of use cases + _reset_everything(random_seed=i, torch_seed=i) + batch = random.choice(range(1, 4)) + heads = random.choice(range(1, 4)) + + # Adjust max seqlens accordingly so we get somewhat more consistent runtimes + max_Q = max_Q_ // batch + max_KV_total = max_KV_total_ // heads + + if backend == "blackwell-fmha": + head_dim_choices = [32, 64, 128] + elif backend == "hopper-fmha": + head_dim_choices = [32, 64, 128] + else: + assert backend == "cutlass-fmha" + head_dim_choices = range(8, 256 + 1, 8) + + head_dim = random.choice(head_dim_choices) + + seqlen_Q = random.choice(range(8, max_Q + 1)) + + num_splits = random.choice(range(2, 9)) + + # Generate split lengths that don't exceed max_KV_total when summed + seqlen_KV_list = [] + remaining = max_KV_total + for j in range(num_splits - 1): + # Each split gets at least 8, up to remaining + max_for_split = max(8, remaining - 8 * (num_splits - j - 1)) + seqlen_KV = random.choice(range(8, min(max_for_split, remaining) + 1)) + seqlen_KV_list.append(seqlen_KV) + remaining -= seqlen_KV + + # Last split gets what's left (at least 8) + seqlen_KV_list.append(max(8, remaining)) + + # torch compile only affects the merge op, and it's slow, so we skew the likelihood + torch_compile = random.choice([True, False, False, False]) + + self._test( + batch=batch, + heads=heads, + head_dim=head_dim, + seqlen_Q=seqlen_Q, + seqlen_KV_list=seqlen_KV_list, + backend=backend, + torch_compile=torch_compile, + ) + + @skip_if_libnatten_is_not_supported() + def test_attention_merge_cutlass_fmha_fast(self): + self._test_randsweep(backend="cutlass-fmha", num_tests=10) + + @skip_if_not_running_extended_tests() + @skip_if_libnatten_is_not_supported() + def test_attention_merge_cutlass_fmha_extended(self): + self._test_randsweep(backend="cutlass-fmha", num_tests=RAND_SWEEP_TESTS) + + @skip_if_hopper_kernels_not_supported() + @skip_if_libnatten_is_not_supported() + def test_attention_merge_hopper_fmha_fast(self): + self._test_randsweep(backend="hopper-fmha", num_tests=10) + + @skip_if_hopper_kernels_not_supported() + @skip_if_not_running_extended_tests() + @skip_if_libnatten_is_not_supported() + def test_attention_merge_hopper_fmha_extended(self): + self._test_randsweep(backend="hopper-fmha", num_tests=RAND_SWEEP_TESTS) + + @skip_if_blackwell_kernels_not_supported() + @skip_if_libnatten_is_not_supported() + def test_attention_merge_blackwell_fmha_fast(self): + self._test_randsweep(backend="blackwell-fmha", num_tests=10) + + @skip_if_blackwell_kernels_not_supported() + @skip_if_not_running_extended_tests() + @skip_if_libnatten_is_not_supported() + def test_attention_merge_blackwell_fmha_extended(self): + self._test_randsweep(backend="blackwell-fmha", num_tests=RAND_SWEEP_TESTS) + + +if __name__ == "__main__": + unittest.main() diff --git a/natten/tests/test_blackwell_fna.py b/natten/tests/test_blackwell_fna.py new file mode 100644 index 00000000..2440ce8d --- /dev/null +++ b/natten/tests/test_blackwell_fna.py @@ -0,0 +1,556 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +import math +import random +import unittest +from itertools import product + +import torch +from natten._environment import _NUM_RAND_SWEEP_TESTS as RAND_SWEEP_TESTS +from natten.backends.configs.cutlass_blackwell import ( + get_all_backward_configs, + get_all_forward_configs, +) +from natten.utils.dtype import is_fp8 +from natten.utils.testing import ( + skip_if_blackwell_kernels_not_supported, + skip_if_libnatten_is_not_supported, + skip_if_not_running_extended_tests, +) + +from .utils import logger, NattenBackendTester + + +def _reset_everything(random_seed: int = 42, torch_seed: int = 42): + from natten.context import ( + NattenContext, + set_memory_usage_preference, + use_kv_parallelism_in_fused_na, + ) + + NattenContext.reset() + set_memory_usage_preference("unrestricted") + use_kv_parallelism_in_fused_na(True) + + random.seed(random_seed) + torch.manual_seed(torch_seed) + logger.debug(f"Reset seeds: {random_seed=}, {torch_seed=}") + torch.cuda.empty_cache() + torch.use_deterministic_algorithms(False) + + +class BlackwellFNABackendTest(unittest.TestCase): + def setUp(self): + _reset_everything() + + def tearDown(self): + _reset_everything() + + def _test_all_dtypes_against_cutlass_2x_fna( + self, + batch, + heads, + head_dim, + input_shape, + kernel_size, + stride, + dilation, + is_causal=None, + configs_to_test=None, + heads_kv=None, + ): + torch.set_default_device("cuda") + assert isinstance(input_shape, tuple) + na_dim = len(input_shape) + assert na_dim in [1, 2, 3], "Only supports NA1D, 2D, 3D." + + tester = NattenBackendTester( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + input_shape=input_shape, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + test_backprop=True, + reference_backend="cutlass-fna", + reference_fmha_backend="cutlass-fmha", + dtype=torch.float16, + ) + + ALLOWED_DTYPES = [ + (torch.float16, (1e-2, (1e-2, 1e-2, 1e-2))), + (torch.bfloat16, (5e-2, (1e-2, 1e-2, 1e-2))), + (torch.float8_e4m3fn, (5e-1, None)), + (torch.float8_e5m2, (9e-1, None)), + ] + + test_id = 0 + for dtype, atol in ALLOWED_DTYPES: + if is_fp8(dtype) and head_dim % 16 != 0: + continue + + dummy = torch.empty( + (batch, *input_shape, heads, head_dim), device="cuda", dtype=dtype + ) + forward_configs = get_all_forward_configs(dummy) + backward_configs = get_all_backward_configs(dummy) + assert len(forward_configs) > 0 + assert len(backward_configs) > 0 + + random.shuffle(forward_configs) + random.shuffle(backward_configs) + + for i in range(max(len(forward_configs), len(backward_configs))): + q_tile_shape, kv_tile_shape = forward_configs[i % len(forward_configs)] + backward_q_tile_shape, backward_kv_tile_shape = backward_configs[ + i % len(backward_configs) + ] + + for persistent in [True, False]: + tester.test( + eps=atol, + dtype=dtype, + target_backend="blackwell-fna", + target_fmha_backend="blackwell-fmha", + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + backward_q_tile_shape=backward_q_tile_shape, + backward_kv_tile_shape=backward_kv_tile_shape, + run_persistent_kernel=persistent, + test_backprop=not is_fp8(dtype), + ) + test_id += 1 + if configs_to_test is not None and test_id > configs_to_test: + return + + @skip_if_libnatten_is_not_supported() + @skip_if_blackwell_kernels_not_supported() + def test_1d_against_cutlass_2x(self): + problem_sizes = [ + (1, 4, 2, 64, (128,), (127,), (1,), (1)), + (1, 4, 2, 48, (128,), (15,), (1,), (1)), + (1, 4, 1, 64, (128,), (15,), (1,), (1)), + (1, 2, 1, 72, (128,), (15,), (1,), (1)), + (1, 2, 2, 64, (128,), (15,), (1,), (1)), + (2, 4, 4, 120, (128,), (63,), (31,), (1)), + (2, 4, 2, 96, (128,), (63,), (31,), (1)), + (2, 4, 1, 128, (128,), (63,), (31,), (1)), + (4, 3, 3, 128, (256,), (255,), (82,), (1)), + (4, 3, 1, 128, (256,), (255,), (82,), (1)), + (1, 1, 1, 128, (32768,), (2048,), (2048,), (1)), + (1, 1, 1, 128, (32768,), (2048,), (256,), (1)), + (1, 1, 1, 128, (32768,), (2048,), (128,), (1)), + (1, 1, 1, 128, (32768,), (2048,), (1,), (1)), + (1, 1, 1, 32, (128,), (3,), (2,), (5)), + (1, 1, 1, 32, (128,), (3,), (1,), (1)), + (1, 1, 1, 32, (128,), (3,), (2,), (10)), + (1, 1, 1, 64, (128,), (8,), (7,), (5)), + (1, 1, 1, 128, (128,), (61,), (33,), (1)), + (1, 1, 1, 32, (125,), (3,), (1,), (1)), + (1, 2, 2, 64, (125,), (15,), (1,), (1)), + (1, 1, 1, 128, (256,), (3,), (2,), (10)), + ] + for i, ( + batch, + heads, + heads_kv, + head_dim, + input_shape, + kernel_size, + stride, + dilation, + ) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for causal in [True, False]: + is_causal = (causal,) + self._test_all_dtypes_against_cutlass_2x_fna( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + input_shape=input_shape, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + ) + + @skip_if_libnatten_is_not_supported() + @skip_if_blackwell_kernels_not_supported() + def test_2d_against_cutlass_2x(self): + problem_sizes = [ + (1, 12, 12, 32, (17, 16), (3, 8), (1, 2), (1, 1)), + (1, 12, 6, 16, (17, 16), (3, 8), (1, 2), (1, 1)), + (1, 12, 4, 32, (17, 16), (3, 8), (1, 2), (1, 1)), + (1, 2, 2, 32, (5, 25), (5, 16), (1, 16), (1, 1)), + (1, 2, 1, 48, (5, 25), (5, 16), (1, 16), (1, 1)), + (1, 1, 1, 32, (84, 19), (7, 3), (1, 1), (1, 1)), + (1, 1, 1, 32, (84, 19), (7, 3), (1, 1), (5, 1)), + (1, 1, 1, 128, (19, 29), (8, 8), (1, 1), (2, 3)), + (1, 1, 1, 96, (48, 17), (24, 16), (1, 1), (2, 1)), + (1, 1, 1, 128, (67, 80), (12, 7), (8, 4), (5, 11)), + (1, 1, 1, 72, (8, 8), (8, 8), (1, 1), (1, 1)), + (1, 1, 1, 128, (33, 33), (24, 16), (1, 1), (1, 1)), + (1, 1, 1, 32, (16, 16), (16, 16), (1, 1), (1, 1)), + (1, 1, 1, 128, (44, 80), (44, 80), (1, 1), (1, 1)), + (1, 1, 1, 32, (40, 20), (3, 7), (1, 1), (1, 1)), + (1, 1, 1, 32, (16, 16), (3, 3), (1, 1), (1, 1)), + (1, 1, 1, 128, (44, 80), (9, 10), (1, 1), (1, 1)), + (1, 1, 1, 64, (28, 40), (17, 31), (1, 1), (1, 1)), + (1, 1, 1, 64, (36, 40), (36, 40), (12, 13), (1, 1)), + (1, 1, 1, 128, (44, 80), (44, 80), (4, 8), (1, 1)), + ] + for i, ( + batch, + heads, + heads_kv, + head_dim, + input_shape, + kernel_size, + stride, + dilation, + ) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for causal_x, causal_y in product([True, False], [True, False]): + is_causal = (causal_x, causal_y) + self._test_all_dtypes_against_cutlass_2x_fna( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + input_shape=input_shape, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + ) + + @skip_if_not_running_extended_tests() + @skip_if_libnatten_is_not_supported() + @skip_if_blackwell_kernels_not_supported() + def test_2d_against_cutlass_2x_extended(self): + problem_sizes = [ + (1, 6, 2, 32, (5, 25), (5, 16), (1, 16), (1, 1)), + (1, 2, 2, 32, (5, 25), (5, 16), (1, 16), (1, 1)), + (1, 2, 1, 32, (5, 25), (5, 16), (4, 16), (1, 1)), + (1, 1, 1, 32, (84, 69), (7, 68), (1, 1), (1, 1)), + (1, 1, 1, 32, (84, 69), (7, 23), (1, 1), (1, 1)), + (1, 1, 1, 32, (84, 69), (7, 20), (1, 6), (1, 1)), + (1, 1, 1, 32, (84, 69), (7, 68), (1, 1), (5, 1)), + (1, 1, 1, 32, (84, 69), (7, 23), (1, 1), (5, 1)), + (1, 1, 1, 32, (84, 69), (7, 20), (1, 6), (5, 1)), + (1, 1, 1, 128, (128, 128), (8, 8), (1, 1), (1, 1)), + (1, 1, 1, 128, (128, 128), (8, 8), (1, 1), (4, 4)), + (1, 1, 1, 128, (64, 64), (32, 32), (1, 1), (2, 2)), + (1, 1, 1, 128, (64, 64), (32, 32), (1, 1), (1, 2)), + (1, 1, 1, 128, (48, 48), (24, 24), (1, 1), (1, 2)), + (1, 1, 1, 128, (17, 48), (16, 24), (1, 1), (1, 2)), + (1, 1, 1, 128, (48, 48), (24, 24), (1, 1), (2, 2)), + (1, 1, 1, 128, (48, 17), (24, 16), (1, 1), (1, 1)), + (1, 1, 1, 128, (72, 80), (24, 16), (1, 1), (3, 5)), + (1, 1, 1, 128, (48, 17), (24, 16), (1, 1), (1, 1)), + (1, 1, 1, 128, (44, 80), (32, 32), (22, 16), (1, 1)), + (1, 1, 1, 128, (44, 80), (24, 24), (1, 1), (1, 1)), + (1, 1, 1, 128, (44, 80), (24, 24), (8, 16), (1, 1)), + (1, 1, 1, 128, (44, 80), (24, 16), (1, 1), (1, 1)), + (1, 1, 1, 128, (44, 80), (24, 16), (8, 16), (1, 1)), + (1, 1, 1, 64, (28, 40), (28, 40), (1, 1), (1, 1)), + (1, 1, 1, 32, (16, 16), (16, 16), (4, 5), (1, 1)), + (1, 1, 1, 32, (16, 16), (16, 16), (4, 8), (1, 1)), + (1, 1, 1, 32, (16, 16), (15, 15), (1, 1), (1, 1)), + (1, 1, 1, 32, (16, 16), (14, 14), (1, 1), (1, 1)), + (1, 1, 1, 64, (36, 40), (36, 40), (1, 1), (1, 1)), + (1, 1, 1, 128, (48, 80), (24, 24), (8, 8), (1, 1)), + ] + for i, ( + batch, + heads, + heads_kv, + head_dim, + input_shape, + kernel_size, + stride, + dilation, + ) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for causal_x, causal_y in product([True, False], [True, False]): + is_causal = (causal_x, causal_y) + self._test_all_dtypes_against_cutlass_2x_fna( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + input_shape=input_shape, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + ) + + @skip_if_libnatten_is_not_supported() + @skip_if_blackwell_kernels_not_supported() + def test_3d_against_cutlass_2x(self): + problem_sizes = [ + (4, 8, 4, 48, (32, 10, 10), (7, 3, 3), (5, 1, 1), (1, 2, 3)), + (4, 8, 2, 64, (32, 10, 10), (7, 3, 3), (5, 1, 1), (1, 2, 3)), + (4, 8, 1, 64, (32, 10, 10), (7, 3, 3), (5, 1, 1), (1, 2, 3)), + (1, 1, 1, 64, (18, 37, 12), (14, 16, 12), (12, 8, 1), (1, 2, 1)), + (1, 1, 1, 32, (13, 11, 9), (3, 4, 3), (2, 3, 3), (1, 1, 1)), + (1, 1, 1, 32, (13, 11, 9), (3, 4, 3), (2, 3, 3), (3, 2, 2)), + (1, 4, 4, 32, (8, 8, 16), (3, 3, 3), (2, 1, 2), (2, 2, 4)), + (1, 1, 1, 64, (18, 37, 12), (14, 16, 12), (12, 8, 1), (1, 2, 1)), + (1, 1, 1, 72, (57, 20, 88), (20, 4, 6), (1, 1, 1), (2, 1, 1)), + (1, 1, 1, 96, (57, 32, 32), (10, 32, 32), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 120, (32, 64, 64), (16, 16, 16), (1, 2, 2), (1, 1, 1)), + (1, 1, 1, 128, (30, 48, 80), (18, 24, 24), (16, 8, 8), (1, 1, 1)), + (1, 1, 1, 128, (16, 44, 80), (12, 32, 32), (8, 22, 16), (1, 1, 1)), + (1, 1, 1, 128, (31, 32, 32), (10, 32, 32), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 128, (30, 32, 32), (10, 32, 32), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 32, (16, 16, 16), (16, 16, 16), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 32, (8, 4, 8), (7, 3, 7), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 32, (16, 16, 16), (15, 15, 15), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 64, (24, 28, 40), (24, 28, 40), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 32, (16, 16, 16), (16, 16, 16), (2, 4, 5), (1, 1, 1)), + ] + for i, ( + batch, + heads, + heads_kv, + head_dim, + input_shape, + kernel_size, + stride, + dilation, + ) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for causal_x, causal_y, causal_z in product( + [True, False], [True, False], [True, False] + ): + is_causal = (causal_x, causal_y, causal_z) + self._test_all_dtypes_against_cutlass_2x_fna( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + input_shape=input_shape, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + ) + + @skip_if_not_running_extended_tests() + @skip_if_libnatten_is_not_supported() + @skip_if_blackwell_kernels_not_supported() + def test_3d_against_cutlass_2x_extended(self): + problem_sizes = [ + (1, 12, 12, 64, (32, 8, 8), (7, 5, 5), (2, 1, 3), (2, 1, 1)), + (1, 12, 6, 64, (32, 8, 8), (7, 5, 5), (2, 1, 3), (2, 1, 1)), + (1, 12, 4, 64, (32, 8, 8), (7, 5, 5), (2, 1, 3), (2, 1, 1)), + (1, 12, 1, 64, (32, 8, 8), (7, 5, 5), (2, 1, 3), (2, 1, 1)), + (2, 2, 2, 32, (8, 8, 10), (3, 4, 3), (3, 4, 1), (1, 1, 1)), + (2, 2, 1, 32, (8, 8, 10), (3, 4, 3), (3, 4, 1), (1, 1, 1)), + (1, 1, 1, 64, (18, 37, 12), (14, 16, 12), (12, 8, 1), (1, 1, 1)), + (1, 1, 1, 32, (13, 11, 9), (3, 4, 3), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 32, (13, 11, 9), (3, 4, 3), (1, 1, 1), (3, 2, 2)), + (1, 1, 1, 32, (8, 8, 4), (3, 4, 3), (1, 1, 1), (1, 1, 1)), + (1, 2, 2, 32, (8, 8, 12), (5, 8, 11), (2, 3, 4), (1, 1, 1)), + (1, 1, 1, 64, (18, 37, 12), (14, 16, 12), (12, 8, 1), (1, 1, 1)), + (2, 3, 3, 32, (5, 45, 73), (2, 7, 32), (1, 4, 32), (2, 2, 2)), + (1, 1, 1, 128, (57, 20, 88), (19, 4, 6), (13, 1, 6), (2, 2, 7)), + (2, 3, 3, 128, (57, 20, 88), (19, 4, 6), (13, 1, 6), (2, 2, 7)), + (2, 3, 1, 128, (57, 20, 88), (19, 4, 6), (13, 1, 6), (2, 2, 7)), + (1, 1, 1, 128, (32, 64, 64), (16, 16, 16), (2, 1, 1), (2, 2, 3)), + (1, 1, 1, 128, (61, 61, 61), (10, 10, 10), (1, 1, 1), (1, 1, 2)), + (1, 1, 1, 128, (61, 61, 61), (10, 10, 10), (1, 1, 1), (1, 2, 1)), + (1, 1, 1, 128, (61, 61, 61), (10, 10, 10), (1, 1, 1), (2, 2, 2)), + (1, 1, 1, 128, (32, 64, 64), (16, 16, 16), (2, 1, 1), (1, 1, 1)), + (1, 1, 1, 128, (32, 64, 64), (16, 16, 16), (1, 1, 2), (1, 1, 1)), + (1, 1, 1, 32, (16, 16, 16), (16, 16, 16), (8, 4, 8), (1, 1, 1)), + (1, 1, 1, 64, (24, 36, 40), (24, 36, 40), (10, 12, 13), (1, 1, 1)), + (1, 1, 1, 128, (16, 44, 80), (16, 44, 80), (8, 4, 8), (1, 1, 1)), + (1, 1, 1, 128, (30, 44, 80), (18, 24, 24), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 128, (30, 44, 80), (18, 24, 24), (2, 8, 16), (1, 1, 1)), + (1, 1, 1, 128, (30, 44, 80), (24, 24, 16), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 128, (30, 44, 80), (24, 24, 16), (2, 8, 16), (1, 1, 1)), + (1, 1, 1, 32, (16, 16, 16), (14, 14, 14), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 32, (16, 16, 16), (3, 3, 3), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 128, (16, 44, 80), (8, 9, 10), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 64, (24, 28, 40), (11, 17, 31), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 64, (24, 36, 40), (24, 36, 40), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 128, (16, 44, 80), (16, 44, 80), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 128, (30, 48, 18), (30, 48, 18), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 128, (8, 8, 8), (8, 8, 8), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 128, (30, 48, 17), (18, 24, 16), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 128, (33, 33, 33), (18, 24, 16), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 128, (30, 48, 18), (3, 2, 8), (1, 1, 1), (8, 18, 2)), + (1, 1, 1, 128, (61, 32, 32), (10, 32, 32), (1, 1, 1), (2, 1, 1)), + (1, 1, 1, 128, (57, 32, 32), (10, 32, 32), (1, 1, 1), (2, 1, 1)), + (1, 1, 1, 128, (57, 20, 88), (20, 8, 16), (1, 1, 1), (2, 1, 1)), + (3, 1, 1, 64, (18, 37, 12), (14, 16, 12), (12, 8, 6), (1, 2, 1)), + (1, 1, 1, 128, (32, 64, 64), (16, 16, 16), (1, 1, 2), (1, 3, 2)), + (1, 1, 1, 128, (48, 64, 64), (7, 15, 11), (1, 2, 2), (5, 3, 2)), + ] + for i, ( + batch, + heads, + heads_kv, + head_dim, + input_shape, + kernel_size, + stride, + dilation, + ) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for causal_x, causal_y, causal_z in product( + [True, False], [True, False], [True, False] + ): + is_causal = (causal_x, causal_y, causal_z) + self._test_all_dtypes_against_cutlass_2x_fna( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + input_shape=input_shape, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + ) + + def _test_randsweep_against_cutlass_2x( + self, na_dim, max_tests=1000, configs_to_test=5 + ): + max_seqlen = 2**17 + # max size per-dim for different profiles + max_size = {1: 2**15, 2: 128, 3: 96} + + # seqlen limit for freely choosing batch and heads + seqlen_limit_batched = 2**13 + + for i in range(max_tests): + # to help with reproducibility of use cases + _reset_everything(random_seed=i, torch_seed=i) + + input_shape = [] + for j in range(na_dim): + input_shape.append(random.choice(range(4, max_size[na_dim] + 1))) + + while math.prod(input_shape) > max_seqlen: + dim_to_cut = random.choice(range(na_dim)) + input_shape[dim_to_cut] = max(4, int(input_shape[dim_to_cut] * 0.1)) + + input_shape = tuple(input_shape) + assert math.prod(input_shape) <= max_seqlen + + max_heads = min( + max(1, (seqlen_limit_batched // math.prod(input_shape)) * 4), 4 + ) + heads = random.choice(range(1, max_heads + 1)) + + max_batch = min( + max(1, ((seqlen_limit_batched // math.prod(input_shape)) * 4) // heads), + 4, + ) + batch = random.choice(range(1, max_batch + 1)) + + heads_kv = random.choice([i for i in range(1, heads + 1) if heads % i == 0]) + head_dim = random.choice(list(range(8, 129, 8))) + + kernel_size = tuple(random.choice(range(2, x + 1)) for x in input_shape) + stride = tuple(random.choice(range(1, k + 1)) for k in kernel_size) + dilation = tuple( + random.choice(range(1, x // k + 1)) + for x, k in zip(input_shape, kernel_size) + ) + is_causal = tuple(random.choice([False, True]) for _ in range(na_dim)) + + self._test_all_dtypes_against_cutlass_2x_fna( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + input_shape=input_shape, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + configs_to_test=configs_to_test, + ) + + @skip_if_libnatten_is_not_supported() + @skip_if_blackwell_kernels_not_supported() + def test_randsweep_1d_against_cutlass_2x_quick(self): + self._test_randsweep_against_cutlass_2x(1, max_tests=10, configs_to_test=None) + + @skip_if_libnatten_is_not_supported() + @skip_if_blackwell_kernels_not_supported() + def test_randsweep_2d_against_cutlass_2x_quick(self): + self._test_randsweep_against_cutlass_2x(2, max_tests=10, configs_to_test=10) + + @skip_if_libnatten_is_not_supported() + @skip_if_blackwell_kernels_not_supported() + def test_randsweep_3d_against_cutlass_2x_quick(self): + self._test_randsweep_against_cutlass_2x(3, max_tests=10, configs_to_test=10) + + @skip_if_not_running_extended_tests() + @skip_if_libnatten_is_not_supported() + @skip_if_blackwell_kernels_not_supported() + def test_randsweep_1d_against_cutlass_2x(self): + self._test_randsweep_against_cutlass_2x(1, max_tests=RAND_SWEEP_TESTS) + + @skip_if_not_running_extended_tests() + @skip_if_libnatten_is_not_supported() + @skip_if_blackwell_kernels_not_supported() + def test_randsweep_2d_against_cutlass_2x(self): + # NOTE: Blackwell FNA has more variations than Hopper FNA due to persistent scheduling being + # separate from the config, 2 additional dtypes (even though they're forward only), and more + # configs in 2D and 3D (despite fewer in 1D). + # This has caused this entire test, with randsweep over 1000 use cases to take close to 6 + # hours. We need to cut down on configs like we do for cutlass-fna to make it run in a more + # reasonable time frame. + self._test_randsweep_against_cutlass_2x( + 2, max_tests=RAND_SWEEP_TESTS, configs_to_test=5 + ) + + @skip_if_not_running_extended_tests() + @skip_if_libnatten_is_not_supported() + @skip_if_blackwell_kernels_not_supported() + def test_randsweep_3d_against_cutlass_2x(self): + # NOTE: Blackwell FNA has more variations than Hopper FNA due to persistent scheduling being + # separate from the config, 2 additional dtypes (even though they're forward only), and more + # configs in 2D and 3D (despite fewer in 1D). + # This has caused this entire test, with randsweep over 1000 use cases to take close to 6 + # hours. We need to cut down on configs like we do for cutlass-fna to make it run in a more + # reasonable time frame. + self._test_randsweep_against_cutlass_2x( + 3, max_tests=RAND_SWEEP_TESTS, configs_to_test=5 + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/natten/tests/test_compute_delta.py b/natten/tests/test_compute_delta.py new file mode 100644 index 00000000..c5312e72 --- /dev/null +++ b/natten/tests/test_compute_delta.py @@ -0,0 +1,125 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +import random +import unittest + +import torch +from natten._libnatten import compute_delta # type: ignore +from natten.utils.testing import ( + skip_if_libnatten_is_not_supported, + supports_bfloat16, + supports_float16, +) + +from .utils import logger + + +def _reset_everything(random_seed: int = 42, torch_seed: int = 42): + random.seed(random_seed) + torch.manual_seed(torch_seed) + logger.debug(f"Reset seeds: {random_seed=}, {torch_seed=}") + torch.use_deterministic_algorithms(True) + + +def compute_delta_pt(out: torch.Tensor, d_out: torch.Tensor, dtype) -> torch.Tensor: + assert out.dim() == d_out.dim() == 4 + assert out.shape == d_out.shape + with torch.no_grad(): + return (out.clone().to(dtype) * d_out.clone().to(dtype)).sum(-1) + + +class ComputeDeltaTests(unittest.TestCase): + def setUp(self): + _reset_everything() + + def tearDown(self): + _reset_everything() + + def _test_against_reference( + self, batch, seqlen, heads, head_dim, eps, dtype, dtype_out + ): + input_shape = (batch, seqlen, heads, head_dim) + out = torch.randn(input_shape, device="cuda", dtype=dtype) + d_out = torch.randn_like(out) + + with torch.no_grad(): + delta_ref = compute_delta_pt(out, d_out, dtype_out) + + delta = torch.empty(input_shape[:-1], device="cuda", dtype=dtype_out) + compute_delta(out, d_out, delta) + + torch.testing.assert_close(delta, delta_ref, atol=eps, rtol=0) + + def _test_all_dtypes_against_reference(self, batch, seqlen, heads, head_dim): + torch.set_default_device("cuda") + self._test_against_reference( + batch=batch, + seqlen=seqlen, + heads=heads, + head_dim=head_dim, + eps=1e-3, + dtype=torch.float32, + dtype_out=torch.float32, + ) + if supports_float16(torch.get_default_device()): + self._test_against_reference( + batch=batch, + seqlen=seqlen, + heads=heads, + head_dim=head_dim, + eps=1e-1, + dtype=torch.float16, + dtype_out=torch.float32, + ) + + if supports_bfloat16(torch.get_default_device()): + self._test_against_reference( + batch=batch, + seqlen=seqlen, + heads=heads, + head_dim=head_dim, + eps=2e-1, + dtype=torch.bfloat16, + dtype_out=torch.float32, + ) + + @skip_if_libnatten_is_not_supported() + def test_against_pt_reference(self): + # (batch, seqlen, heads, head_dim) + input_sizes = [ + (1, 64, 2, 128), + (128, 56, 4, 2), + (128, 56, 8, 16), + (1, 1235, 56, 512), + (2, 128, 4, 48), + (5, 50, 1, 60), + (1, 127, 1, 256), + ] + for i, (batch, seqlen, heads, head_dim) in enumerate(input_sizes): + _reset_everything(random_seed=i, torch_seed=i) + self._test_all_dtypes_against_reference(batch, seqlen, heads, head_dim) + + +if __name__ == "__main__": + unittest.main() diff --git a/natten/tests/test_fmha.py b/natten/tests/test_fmha.py new file mode 100644 index 00000000..d12ad60f --- /dev/null +++ b/natten/tests/test_fmha.py @@ -0,0 +1,1449 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +import random +import unittest +from typing import Optional, Tuple, Union + +import natten # noqa: F401 +import pytest +import torch +from natten._environment import _NUM_RAND_SWEEP_TESTS as RAND_SWEEP_TESTS +from natten.backends.configs.cutlass import ( + get_all_fmha_backward_configs, + get_all_fmha_forward_configs, +) +from natten.backends.configs.cutlass_blackwell import ( + get_all_fmha_backward_configs as get_all_blackwell_fmha_backward_configs, + get_all_fmha_forward_configs as get_all_blackwell_fmha_forward_configs, +) +from natten.backends.configs.cutlass_hopper import ( + get_all_fmha_backward_configs as get_all_hopper_fmha_backward_configs, + get_all_fmha_forward_configs as get_all_hopper_fmha_forward_configs, +) +from natten.functional import attention +from natten.types import KernelSchedule +from natten.utils.dtype import is_fp8 +from natten.utils.testing import ( + skip_if_blackwell_kernels_not_supported, + skip_if_hopper_kernels_not_supported, + skip_if_libnatten_is_not_supported, + skip_if_not_running_extended_tests, + supports_bfloat16, + supports_float16, +) +from torch import Tensor + +from .utils import logger + + +def _reset_everything(random_seed: int = 42, torch_seed: int = 42): + from natten.context import ( + NattenContext, + set_memory_usage_preference, + use_kv_parallelism_in_fused_na, + ) + + NattenContext.reset() + set_memory_usage_preference("unrestricted") + use_kv_parallelism_in_fused_na(True) + + random.seed(random_seed) + torch.manual_seed(torch_seed) + logger.debug(f"Reset seeds: {random_seed=}, {torch_seed=}") + torch.cuda.empty_cache() + torch.use_deterministic_algorithms(False) + + +def compute_sdpa_reference( + batch, + heads, + heads_kv, + dim, + dim_v, + seqlen_q, + seqlen_kv, + is_causal, + dtype=torch.float32, +): + dim_v = dim_v or dim + assert heads % heads_kv == 0 + h_k = heads // heads_kv + with torch.no_grad(): + q, k, v, d_out = ( + torch.randn((batch, heads, seqlen_q, dim), device="cuda", dtype=dtype), + torch.randn((batch, heads_kv, seqlen_kv, dim), device="cuda", dtype=dtype), + torch.randn( + (batch, heads_kv, seqlen_kv, dim_v), device="cuda", dtype=dtype + ), + torch.randn((batch, heads, seqlen_q, dim_v), device="cuda", dtype=dtype) + * 0.05, + ) + q_, k_, v_, d_out_ = ( + q.clone().permute(0, 2, 1, 3).contiguous(), + k.clone().permute(0, 2, 1, 3).contiguous(), + v.clone().permute(0, 2, 1, 3).contiguous(), + d_out.clone().permute(0, 2, 1, 3).contiguous(), + ) + + q = q.requires_grad_(True) + k = k.requires_grad_(True) + v = v.requires_grad_(True) + d_out = d_out.requires_grad_(True) + + k_final, v_final = k, v + # Decomposed GQA/MQA implementation for torch SDPA via explicit repeats + if h_k > 1: + k_final = torch.repeat_interleave(k, repeats=h_k, dim=1, output_size=heads) + v_final = torch.repeat_interleave(v, repeats=h_k, dim=1, output_size=heads) + + assert k_final.shape[:2] == q.shape[:2] + assert v_final.shape[:2] == q.shape[:2] + assert k_final.shape[-1] == q.shape[-1] + assert v_final.shape[-1] == q.shape[-1] + assert k_final.shape[2] == seqlen_kv + assert v_final.shape[2] == seqlen_kv + + with torch.nn.attention.sdpa_kernel( + backends=[torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION] + ): + out = torch.nn.functional.scaled_dot_product_attention( + q, k_final, v_final, is_causal=is_causal + ) + + lse_ref = None + + with torch.no_grad(): + out_ref = out.clone().permute(0, 2, 1, 3).contiguous().float() + + out.backward(d_out) + with torch.no_grad(): + dq_ref, dk_ref, dv_ref = ( + q.grad.clone().permute(0, 2, 1, 3).contiguous().float(), + k.grad.clone().permute(0, 2, 1, 3).contiguous().float(), + v.grad.clone().permute(0, 2, 1, 3).contiguous().float(), + ) + return (q_, k_, v_, d_out_), (out_ref, lse_ref, dq_ref, dk_ref, dv_ref) + + +def compute_natten_fmha_reference( + batch, + heads, + dim, + dim_v, + seqlen_q, + seqlen_kv, + is_causal, + dtype=torch.float32, + backend="cutlass-fmha", + heads_kv=None, + q_tile_size=None, + kv_tile_size=None, + backward_q_tile_size=None, + backward_kv_tile_size=None, + backward_kv_splits=None, + backward_use_pt_reduction=False, +): + dim_v = dim_v or dim + heads_kv = heads_kv or heads + with torch.no_grad(): + q, k, v, d_out = ( + torch.randn((batch, seqlen_q, heads, dim), device="cuda", dtype=dtype), + torch.randn((batch, seqlen_kv, heads_kv, dim), device="cuda", dtype=dtype), + torch.randn( + (batch, seqlen_kv, heads_kv, dim_v), device="cuda", dtype=dtype + ), + torch.randn((batch, seqlen_q, heads, dim_v), device="cuda", dtype=dtype) + * 0.05, + ) + q_, k_, v_, d_out_ = ( + q.clone(), + k.clone(), + v.clone(), + d_out.clone(), + ) + + q = q.requires_grad_(True) + k = k.requires_grad_(True) + v = v.requires_grad_(True) + d_out = d_out.requires_grad_(True) + + out, lse = attention( + q, + k, + v, + is_causal=is_causal, + backend=backend, + q_tile_size=q_tile_size, + kv_tile_size=kv_tile_size, + backward_q_tile_size=backward_q_tile_size, + backward_kv_tile_size=backward_kv_tile_size, + backward_kv_splits=backward_kv_splits, + backward_use_pt_reduction=backward_use_pt_reduction, + return_lse=True, + ) + + with torch.no_grad(): + out_ref = out.clone().float() + lse_ref = lse.clone().float() + + out.backward(d_out) + with torch.no_grad(): + dq_ref, dk_ref, dv_ref = ( + q.grad.clone().float(), + k.grad.clone().float(), + v.grad.clone().float(), + ) + return (q_, k_, v_, d_out_), (out_ref, lse_ref, dq_ref, dk_ref, dv_ref) + + +# TODO: write a class like the FNA tests +class FMHABackendTest(unittest.TestCase): + def setUp(self): + _reset_everything() + + def tearDown(self): + _reset_everything() + + def _test_against_reference_inputs( + self, + inputs, + reference, + batch: int, + heads: int, + head_dim: int, + seqlen_q: int, + seqlen_kv: int, + is_causal: bool, + atol: Union[float, Tuple[float, float]], + dtype: torch.dtype, + test_backprop: bool, + backend: str, + q_tile_size: Optional[int] = None, + kv_tile_size: Optional[int] = None, + backward_q_tile_size: Optional[int] = None, + backward_kv_tile_size: Optional[int] = None, + backward_kv_splits: Optional[int] = None, + backward_use_pt_reduction: bool = False, + run_persistent_kernel: bool = True, + kernel_schedule: Optional[KernelSchedule] = None, + test_lse: bool = False, + reference_str: str = "unspecified reference", + heads_kv: Optional[int] = None, + head_dim_v: Optional[int] = None, + ): + heads_kv = heads_kv or heads + head_dim_v = head_dim_v or head_dim + determinism = torch.are_deterministic_algorithms_enabled() + logger.debug( + f"Testing FMHA ({backend}) vs {reference_str}:\n" + f"{batch=}, {heads=}, {heads_kv=}, {head_dim=}, {head_dim_v=},\n" + f"{seqlen_q=}, {seqlen_kv=}, {is_causal=}, {dtype=},\n" + f"{q_tile_size=}, {kv_tile_size=}, {kernel_schedule=}, {run_persistent_kernel=}" + + ( + f",\n({determinism=}) {backward_q_tile_size=}, {backward_kv_tile_size=}, " + f"{backward_kv_splits=}, {backward_use_pt_reduction=}." + if test_backprop + else "." + ) + ) + + q, k, v, d_out = inputs + out_ref, lse_ref, dq_ref, dk_ref, dv_ref = reference + q = q.to(dtype) + k = k.to(dtype) + v = v.to(dtype) + d_out = d_out.to(dtype) + + # Run target + q.requires_grad_(test_backprop) + k.requires_grad_(test_backprop) + v.requires_grad_(test_backprop) + d_out.requires_grad_(test_backprop) + + outputs = attention( + q, + k, + v, + is_causal=is_causal, + backend=backend, + q_tile_size=q_tile_size, + kv_tile_size=kv_tile_size, + backward_q_tile_size=backward_q_tile_size, + backward_kv_tile_size=backward_kv_tile_size, + backward_kv_splits=backward_kv_splits, + backward_use_pt_reduction=backward_use_pt_reduction, + run_persistent_kernel=run_persistent_kernel, + kernel_schedule=kernel_schedule, + return_lse=test_lse, + ) + + if test_lse: + assert ( + isinstance(outputs, tuple) + and len(outputs) == 2 + and all(isinstance(x, Tensor) for x in outputs) + ) + out_, lse_ = outputs + + out = out_.data.clone().float() + lse = lse_.clone().float() + else: + assert isinstance(outputs, Tensor) + out_ = outputs + + out = out_.data.clone().float() + + if test_backprop: + dq, dk, dv = None, None, None + out_.backward(d_out) + with torch.no_grad(): + dq, dk, dv = ( + q.grad.clone().float(), + k.grad.clone().float(), + v.grad.clone().float(), + ) + + if isinstance(atol, tuple): + atol_forward, atol_backward = atol + else: + atol_forward, atol_backward = atol, atol + + torch.testing.assert_close(out, out_ref, atol=atol_forward, rtol=0) + if test_lse: + assert ( + lse_ref is not None + ), "Reference did not return LSE. If reference is PyTorch SDPA, it does not have an API for returning LSE. Use CUTLASS FMHA instead!" + + torch.testing.assert_close(lse, lse_ref, atol=atol_forward, rtol=0) + + if test_backprop: + if isinstance(atol_backward, tuple): + assert len(atol_backward) == 3 + atol_dq, atol_dk, atol_dv = atol_backward + else: + assert isinstance(atol_backward, float) or isinstance( + atol_backward, int + ) + atol_dq, atol_dk, atol_dv = atol_backward, atol_backward, atol_backward + torch.testing.assert_close(dq, dq_ref, atol=atol_dq, rtol=0) + torch.testing.assert_close(dk, dk_ref, atol=atol_dk, rtol=0) + torch.testing.assert_close(dv, dv_ref, atol=atol_dv, rtol=0) + + def _test_against_torch_sdpa( + self, + batch: int, + heads: int, + head_dim: int, + seqlen_q: int, + seqlen_kv: int, + is_causal: bool, + atol: Union[float, Tuple[float, float]], + dtype: torch.dtype, + test_backprop: bool, + backend: str, + q_tile_size: Optional[int] = None, + kv_tile_size: Optional[int] = None, + backward_q_tile_size: Optional[int] = None, + backward_kv_tile_size: Optional[int] = None, + backward_kv_splits: Optional[int] = None, + backward_use_pt_reduction: bool = False, + run_persistent_kernel: bool = True, + kernel_schedule: Optional[KernelSchedule] = None, + heads_kv: Optional[int] = None, + head_dim_v: Optional[int] = None, + ): + heads_kv = heads_kv or heads + head_dim_v = head_dim_v or head_dim + sdpa_dtype = dtype if not is_fp8(dtype) else torch.float16 + inputs, reference = compute_sdpa_reference( + batch, + heads, + heads_kv, + head_dim, + head_dim_v, + seqlen_q, + seqlen_kv, + is_causal=is_causal, + dtype=sdpa_dtype, + ) + self._test_against_reference_inputs( + inputs=inputs, + reference=reference, + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + head_dim_v=head_dim_v, + seqlen_q=seqlen_q, + seqlen_kv=seqlen_kv, + is_causal=is_causal, + atol=atol, + dtype=dtype, + test_backprop=test_backprop, + backend=backend, + q_tile_size=q_tile_size, + kv_tile_size=kv_tile_size, + backward_q_tile_size=backward_q_tile_size, + backward_kv_tile_size=backward_kv_tile_size, + backward_kv_splits=backward_kv_splits, + backward_use_pt_reduction=backward_use_pt_reduction, + run_persistent_kernel=run_persistent_kernel, + kernel_schedule=kernel_schedule, + reference_str="torch sdpa", + ) + + def _test_against_natten_fmha( + self, + batch: int, + heads: int, + head_dim: int, + seqlen_q: int, + seqlen_kv: int, + is_causal: bool, + atol: Union[float, Tuple[float, float]], + dtype: torch.dtype, + test_backprop: bool, + test_lse: bool, + backend: str, + q_tile_size: Optional[int] = None, + kv_tile_size: Optional[int] = None, + backward_q_tile_size: Optional[int] = None, + backward_kv_tile_size: Optional[int] = None, + backward_kv_splits: Optional[int] = None, + backward_use_pt_reduction: bool = False, + run_persistent_kernel: bool = True, + kernel_schedule: Optional[KernelSchedule] = None, + head_dim_v: Optional[int] = None, + heads_kv: Optional[int] = None, + reference_backend: str = "cutlass-fmha", + reference_q_tile_size: Optional[int] = None, + reference_kv_tile_size: Optional[int] = None, + reference_backward_q_tile_size: Optional[int] = None, + reference_backward_kv_tile_size: Optional[int] = None, + reference_backward_kv_splits: Optional[int] = None, + reference_backward_use_pt_reduction: bool = False, + ): + head_dim_v = head_dim_v or head_dim + heads_kv = heads_kv or heads + inputs, reference = compute_natten_fmha_reference( + batch, + heads, + head_dim, + head_dim_v, + seqlen_q, + seqlen_kv, + is_causal=is_causal, + dtype=dtype, + backend=reference_backend, + heads_kv=heads_kv, + q_tile_size=reference_q_tile_size, + kv_tile_size=reference_kv_tile_size, + backward_q_tile_size=reference_backward_q_tile_size, + backward_kv_tile_size=reference_backward_kv_tile_size, + backward_kv_splits=reference_backward_kv_splits, + backward_use_pt_reduction=reference_backward_use_pt_reduction, + ) + self._test_against_reference_inputs( + inputs=inputs, + reference=reference, + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + head_dim_v=head_dim_v, + seqlen_q=seqlen_q, + seqlen_kv=seqlen_kv, + is_causal=is_causal, + atol=atol, + dtype=dtype, + test_backprop=test_backprop, + backend=backend, + q_tile_size=q_tile_size, + kv_tile_size=kv_tile_size, + backward_q_tile_size=backward_q_tile_size, + backward_kv_tile_size=backward_kv_tile_size, + backward_kv_splits=backward_kv_splits, + backward_use_pt_reduction=backward_use_pt_reduction, + run_persistent_kernel=run_persistent_kernel, + kernel_schedule=kernel_schedule, + test_lse=test_lse, + reference_str=reference_backend, + ) + + def _test_cutlass_fmha_against_torch_sdpa( + self, + batch, + heads, + head_dim, + head_dim_v, + seqlen_q, + seqlen_kv, + is_causal, + max_configs=10, + ): + torch.set_default_device("cuda") + + is_mla = head_dim_v is not None and head_dim != head_dim_v + ALLOWED_DTYPES = [ + (torch.float32, (1e-4, (1e-2, 1e-4, 1e-4))), + ] + + if supports_float16(torch.get_default_device()): + ALLOWED_DTYPES.append((torch.float16, (1e-2, (1e-2, 1e-2, 1e-2)))) + + if supports_bfloat16(torch.get_default_device()): + ALLOWED_DTYPES.append( + ( + torch.bfloat16, + (5e-2, (2e-2, 5e-2, 5e-2) if is_mla else (1e-2, 1e-2, 1e-2)), + ) + ) + + for dtype, atol in ALLOWED_DTYPES: + + dummy_fwd = torch.randn( + (batch, seqlen_q, heads, max(head_dim, head_dim_v)), + device="cuda", + dtype=dtype, + ) + + dummy_bwd = torch.randn( + (batch, seqlen_kv, heads, max(head_dim, head_dim_v)), + device="cuda", + dtype=dtype, + ) + + forward_configs = get_all_fmha_forward_configs(dummy_fwd) + backward_configs = get_all_fmha_backward_configs(dummy_bwd) + assert len(forward_configs) > 0 + assert len(backward_configs) > 0 + + random.shuffle(forward_configs) + random.shuffle(backward_configs) + + n_configs_to_test = min( + max_configs, max(len(forward_configs), len(backward_configs)) + ) + + for i in range(n_configs_to_test): + q_tile_size, kv_tile_size = forward_configs[i % len(forward_configs)] + backward_q_tile_size, backward_kv_tile_size = backward_configs[ + i % len(backward_configs) + ] + + self._test_against_torch_sdpa( + batch=batch, + heads=heads, + head_dim=head_dim, + head_dim_v=head_dim_v, + seqlen_q=seqlen_q, + seqlen_kv=seqlen_kv, + is_causal=is_causal, + dtype=dtype, + test_backprop=True, + atol=atol, + backend="cutlass-fmha", + q_tile_size=q_tile_size, + kv_tile_size=kv_tile_size, + backward_q_tile_size=backward_q_tile_size, + backward_kv_tile_size=backward_kv_tile_size, + ) + + def _test_cutlass_blackwell_fmha_against_torch_sdpa( + self, + batch, + heads, + head_dim, + seqlen_q, + seqlen_kv, + is_causal, + heads_kv=None, + ): + torch.set_default_device("cuda") + + is_gqa = heads_kv is not None and heads != heads_kv + ALLOWED_DTYPES = [ + (torch.float16, (1e-2, (1e-2, 1e-2, 1e-2))), + ( + torch.bfloat16, + (5e-2, (1e-2, 1e-2, 1e-2) if not is_gqa else (1e-2, 5e-2, 5e-2)), + ), + (torch.float8_e4m3fn, (5e-1, None)), + (torch.float8_e5m2, (9e-1, None)), + ] + + for dtype, atol in ALLOWED_DTYPES: + if is_fp8(dtype) and head_dim % 16 != 0: + continue + + dummy = torch.empty( + (batch, seqlen_kv, heads, head_dim), device="cuda", dtype=dtype + ) + + forward_configs = get_all_blackwell_fmha_forward_configs(dummy) + backward_configs = get_all_blackwell_fmha_backward_configs(dummy) + assert len(forward_configs) > 0 + assert len(backward_configs) > 0 + + random.shuffle(forward_configs) + random.shuffle(backward_configs) + + for i in range(max(len(forward_configs), len(backward_configs))): + q_tile_size, kv_tile_size = forward_configs[i] + backward_q_tile_size, backward_kv_tile_size = backward_configs[ + i % len(backward_configs) + ] + + self._test_against_torch_sdpa( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + seqlen_q=seqlen_q, + seqlen_kv=seqlen_kv, + is_causal=is_causal, + dtype=dtype, + test_backprop=not is_fp8(dtype), + atol=atol, + backend="blackwell-fmha", + q_tile_size=q_tile_size, + kv_tile_size=kv_tile_size, + backward_q_tile_size=backward_q_tile_size, + backward_kv_tile_size=backward_kv_tile_size, + ) + + def _test_cutlass_hopper_fmha_against_torch_sdpa( + self, + batch, + heads, + head_dim, + seqlen_q, + seqlen_kv, + is_causal, + ): + torch.set_default_device("cuda") + + ALLOWED_DTYPES = [ + (torch.float16, (1e-2, (1e-2, 1e-2, 1e-2))), + (torch.bfloat16, (5e-2, (1e-2, 2e-2, 2e-2))), + ] + + for dtype, atol in ALLOWED_DTYPES: + + dummy = torch.randn( + (batch, seqlen_kv, heads, head_dim), device="cuda", dtype=dtype + ) + + forward_configs = get_all_hopper_fmha_forward_configs(dummy) + backward_configs = get_all_hopper_fmha_backward_configs(dummy) + assert len(forward_configs) > 0 + assert len(backward_configs) > 0 + + random.shuffle(forward_configs) + random.shuffle(backward_configs) + + for i in range(max(len(forward_configs), len(backward_configs))): + (q_tile_size, kv_tile_size), schedule = forward_configs[ + i % len(forward_configs) + ] + backward_q_tile_size, backward_kv_tile_size = backward_configs[ + i % len(backward_configs) + ] + + self._test_against_torch_sdpa( + batch=batch, + heads=heads, + head_dim=head_dim, + seqlen_q=seqlen_q, + seqlen_kv=seqlen_kv, + is_causal=is_causal, + dtype=dtype, + test_backprop=True, + atol=atol, + backend="hopper-fmha", + q_tile_size=q_tile_size, + kv_tile_size=kv_tile_size, + kernel_schedule=schedule, + backward_q_tile_size=backward_q_tile_size, + backward_kv_tile_size=backward_kv_tile_size, + ) + + def _test_cutlass_hopper_fmha_against_natten( + self, + batch, + heads, + head_dim, + seqlen_q, + seqlen_kv, + is_causal, + ): + torch.set_default_device("cuda") + + ALLOWED_DTYPES = [ + (torch.float16, (1e-2, (1e-2, 1e-2, 1e-2))), + (torch.bfloat16, (5e-2, (1e-2, 1e-2, 1e-2))), + ] + + for dtype, atol in ALLOWED_DTYPES: + + dummy = torch.randn( + (batch, seqlen_kv, heads, head_dim), device="cuda", dtype=dtype + ) + + forward_configs = get_all_hopper_fmha_forward_configs(dummy) + backward_configs = get_all_hopper_fmha_backward_configs(dummy) + assert len(forward_configs) > 0 + assert len(backward_configs) > 0 + + random.shuffle(forward_configs) + random.shuffle(backward_configs) + + for i in range(max(len(forward_configs), len(backward_configs))): + (q_tile_size, kv_tile_size), schedule = forward_configs[ + i % len(forward_configs) + ] + backward_q_tile_size, backward_kv_tile_size = backward_configs[ + i % len(backward_configs) + ] + + self._test_against_natten_fmha( + batch=batch, + heads=heads, + head_dim=head_dim, + seqlen_q=seqlen_q, + seqlen_kv=seqlen_kv, + is_causal=is_causal, + dtype=dtype, + test_backprop=True, + test_lse=True, + atol=atol, + backend="hopper-fmha", + q_tile_size=q_tile_size, + kv_tile_size=kv_tile_size, + kernel_schedule=schedule, + backward_q_tile_size=backward_q_tile_size, + backward_kv_tile_size=backward_kv_tile_size, + ) + + def _test_cutlass_fmha_determinism( + self, + batch, + heads, + head_dim, + seqlen_q, + seqlen_kv, + is_causal, + max_configs=5, + fail_case=False, + ): + torch.set_default_device("cuda") + + torch.use_deterministic_algorithms(not fail_case) + ALLOWED_DTYPES = [ + (torch.float32, (0, 0)), + (torch.float16, (0, 0)), + (torch.bfloat16, (0, 0)), + ] + + for dtype, atol in ALLOWED_DTYPES: + dummy = torch.empty( + (batch, seqlen_kv, heads, head_dim), device="cuda", dtype=dtype + ) + + forward_configs = get_all_fmha_forward_configs(dummy) + backward_configs = get_all_fmha_backward_configs(dummy) + assert len(forward_configs) > 0 + assert len(backward_configs) > 0 + + random.shuffle(forward_configs) + random.shuffle(backward_configs) + + n_configs_to_test = min( + max_configs, max(len(forward_configs), len(backward_configs)) + ) + + for i in range(n_configs_to_test): + q_tile_size, kv_tile_size = forward_configs[i % len(forward_configs)] + backward_q_tile_size, backward_kv_tile_size = backward_configs[ + i % len(backward_configs) + ] + + self._test_against_natten_fmha( + batch=batch, + heads=heads, + head_dim=head_dim, + seqlen_q=seqlen_q, + seqlen_kv=seqlen_kv, + is_causal=is_causal, + dtype=dtype, + test_backprop=True, + test_lse=True, + atol=atol, + backend="cutlass-fmha", + q_tile_size=q_tile_size, + kv_tile_size=kv_tile_size, + backward_q_tile_size=backward_q_tile_size, + backward_kv_tile_size=backward_kv_tile_size, + backward_kv_splits=1, + backward_use_pt_reduction=False, + reference_backend="cutlass-fmha", + reference_q_tile_size=q_tile_size, + reference_kv_tile_size=kv_tile_size, + reference_backward_q_tile_size=backward_q_tile_size, + reference_backward_kv_tile_size=backward_kv_tile_size, + reference_backward_kv_splits=1, + reference_backward_use_pt_reduction=False, + ) + + torch.use_deterministic_algorithms(False) + + def _test_cutlass_blackwell_fmha_against_natten( + self, + batch, + heads, + head_dim, + seqlen_q, + seqlen_kv, + is_causal, + heads_kv=None, + determinism=False, + fail_case=False, + ): + torch.set_default_device("cuda") + + ALLOWED_DTYPES = [ + (torch.float16, (1e-2, (1e-2, 1e-2, 1e-2))), + (torch.bfloat16, (5e-2, (1e-2, 1e-2, 1e-2))), + ] + + if determinism: + ALLOWED_DTYPES = [ + (torch.float16, (0, 0)), + (torch.bfloat16, (0, 0)), + ] + torch.use_deterministic_algorithms(not fail_case) + + for dtype, atol in ALLOWED_DTYPES: + + dummy = torch.randn( + (batch, seqlen_kv, heads, head_dim), device="cuda", dtype=dtype + ) + + forward_configs = get_all_blackwell_fmha_forward_configs(dummy) + backward_configs = get_all_blackwell_fmha_backward_configs(dummy) + assert len(forward_configs) > 0 + assert len(backward_configs) > 0 + + random.shuffle(forward_configs) + random.shuffle(backward_configs) + + for i in range(max(len(forward_configs), len(backward_configs))): + q_tile_size, kv_tile_size = forward_configs[i % len(forward_configs)] + backward_q_tile_size, backward_kv_tile_size = backward_configs[ + i % len(backward_configs) + ] + + reference_kwargs = {} + if determinism: + reference_kwargs["reference_backend"] = "blackwell-fmha" + reference_kwargs["reference_q_tile_size"] = q_tile_size + reference_kwargs["reference_kv_tile_size"] = kv_tile_size + reference_kwargs["reference_backward_q_tile_size"] = ( + backward_q_tile_size + ) + reference_kwargs["reference_backward_kv_tile_size"] = ( + backward_kv_tile_size + ) + + self._test_against_natten_fmha( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + seqlen_q=seqlen_q, + seqlen_kv=seqlen_kv, + is_causal=is_causal, + dtype=dtype, + test_backprop=True, + test_lse=True, + atol=atol, + backend="blackwell-fmha", + q_tile_size=q_tile_size, + kv_tile_size=kv_tile_size, + backward_q_tile_size=backward_q_tile_size, + backward_kv_tile_size=backward_kv_tile_size, + **reference_kwargs, + ) + + torch.use_deterministic_algorithms(False) + + def _test_backend_against_torch_sdpa( + self, + batch, + heads, + head_dim, + seqlen_q, + seqlen_kv, + is_causal, + backend, + heads_kv=None, + head_dim_v=None, + ): + if backend == "cutlass-fmha": + assert heads_kv is None or heads_kv == heads + return self._test_cutlass_fmha_against_torch_sdpa( + batch=batch, + heads=heads, + head_dim=head_dim, + head_dim_v=head_dim_v, + seqlen_q=seqlen_q, + seqlen_kv=seqlen_kv, + is_causal=is_causal, + ) + elif backend == "blackwell-fmha": + assert ( + head_dim_v is None or head_dim_v == head_dim + ), "Blackwell FMHA does not allow head_dim_v." + + return self._test_cutlass_blackwell_fmha_against_torch_sdpa( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + seqlen_q=seqlen_q, + seqlen_kv=seqlen_kv, + is_causal=is_causal, + ) + elif backend == "hopper-fmha": + assert heads_kv is None or heads_kv == heads + assert ( + head_dim_v is None or head_dim_v == head_dim + ), "Hopper FMHA does not allow head_dim_v." + + return self._test_cutlass_hopper_fmha_against_torch_sdpa( + batch=batch, + heads=heads, + head_dim=head_dim, + seqlen_q=seqlen_q, + seqlen_kv=seqlen_kv, + is_causal=is_causal, + ) + else: + raise NotImplementedError(f"Add {backend=} to tests.") + + # Intended to test LSE and determinism + def _test_backend_against_natten( + self, + batch, + heads, + head_dim, + seqlen_q, + seqlen_kv, + is_causal, + backend, + heads_kv=None, + head_dim_v=None, + determinism: bool = False, + fail_case: bool = False, + ): + if backend == "hopper-fmha": + assert heads_kv is None or heads_kv == heads + assert ( + head_dim_v is None or head_dim_v == head_dim + ), "Hopper FMHA does not allow head_dim_v." + assert not determinism + assert not fail_case + return self._test_cutlass_hopper_fmha_against_natten( + batch=batch, + heads=heads, + head_dim=head_dim, + seqlen_q=seqlen_q, + seqlen_kv=seqlen_kv, + is_causal=is_causal, + ) + elif backend == "cutlass-fmha": + assert heads_kv is None or heads_kv == heads + assert determinism + return self._test_cutlass_fmha_determinism( + batch=batch, + heads=heads, + head_dim=head_dim, + seqlen_q=seqlen_q, + seqlen_kv=seqlen_kv, + is_causal=is_causal, + fail_case=fail_case, + ) + elif backend == "blackwell-fmha": + assert ( + head_dim_v is None or head_dim_v == head_dim + ), "Blackwell FMHA does not allow head_dim_v." + + return self._test_cutlass_blackwell_fmha_against_natten( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + seqlen_q=seqlen_q, + seqlen_kv=seqlen_kv, + is_causal=is_causal, + determinism=determinism, + fail_case=fail_case, + ) + else: + raise NotImplementedError(f"Add {backend=} to tests.") + + def _test_randsweep( + self, + backend, + max_tests=1000, + reference: str = "sdpa", + determinism: bool = False, + ): + assert reference in ("sdpa", "natten") + max_qk = 2**21 + for i in range(max_tests): + # to help with reproducibility of use cases + _reset_everything(random_seed=i, torch_seed=i) + batch = random.choice(range(1, 4)) + + supports_dim_v = False + supports_gqa_mqa = False + if backend == "blackwell-fmha": + head_dim_choices = list(range(8, 129, 8)) + heads_choices = range(1, 8 + 1) + supports_gqa_mqa = True + elif backend == "hopper-fmha": + head_dim_choices = [32, 64, 128] + heads_choices = range(1, 4) + else: + assert backend == "cutlass-fmha" + head_dim_choices = range(8, 1024 + 1, 8) # type: ignore[assignment] + heads_choices = range(1, 4) + supports_dim_v = True + + heads = random.choice(heads_choices) + heads_kv = ( + heads + if not supports_gqa_mqa + else random.choice( + [1] + [i for i in range(1, heads + 1) if heads % i == 0] + ) + ) + assert heads >= heads_kv and heads % heads_kv == 0 + + head_dim = random.choice(head_dim_choices) + head_dim_v = ( + head_dim if not supports_dim_v else random.choice(head_dim_choices) + ) + + seqlen_q = random.choice(range(8, 2**14, 1)) + seqlen_kv = random.choice(range(8, 2**14, 1)) + + while seqlen_q * seqlen_kv > max_qk: + cut_kv = random.choice([True, False]) + if cut_kv: + seqlen_kv = int(seqlen_kv * 0.1) + else: + seqlen_q = int(seqlen_q * 0.1) + + for is_causal in [False, True]: + if reference == "sdpa": + assert not determinism + self._test_backend_against_torch_sdpa( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + head_dim_v=head_dim_v, + seqlen_q=seqlen_q, + seqlen_kv=seqlen_kv, + is_causal=is_causal, + backend=backend, + ) + elif reference == "natten": + self._test_backend_against_natten( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + head_dim_v=head_dim_v, + seqlen_q=seqlen_q, + seqlen_kv=seqlen_kv, + is_causal=is_causal, + backend=backend, + determinism=determinism, + ) + + @skip_if_libnatten_is_not_supported() + def test_cutlass_fmha_fast_against_torch_sdpa(self): + problem_sizes = [ + (1, 1, 32, 64, 128, 128), + (2, 2, 64, 32, 128, 128), + (1, 1, 32, 16, 32, 32), + (1, 1, 32, 48, 128, 128), + (1, 1, 192, 64, 3584, 381), + (1, 1, 288, 256, 12072, 1680), + (1, 1, 32, 64, 128, 128), + (1, 1, 32, 128, 128, 4096), + # + (2, 2, 64, 64, 128, 128), + (1, 1, 32, 32, 32, 32), + (1, 1, 32, 32, 128, 128), + (1, 1, 192, 192, 3584, 381), + (1, 1, 288, 288, 12072, 1680), + (1, 1, 32, 32, 128, 128), + (1, 1, 32, 32, 128, 4096), + (1, 1, 32, 32, 128, 258), + (1, 2, 64, 64, 128, 15), + (1, 1, 32, 32, 8, 17), + (1, 1, 64, 64, 17, 49), + (2, 4, 16, 16, 128, 237), + (4, 3, 48, 48, 256, 33), + (1, 1, 128, 128, 128, 75), + (1, 1, 32, 32, 125, 444), + (1, 2, 64, 64, 125, 231), + (1, 1, 128, 128, 256, 10240), + ] + for i, ( + batch, + heads, + head_dim, + head_dim_v, + seqlen_q, + seqlen_kv, + ) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for is_causal in [False, True]: + self._test_backend_against_torch_sdpa( + batch=batch, + heads=heads, + head_dim=head_dim, + head_dim_v=head_dim_v, + seqlen_q=seqlen_q, + seqlen_kv=seqlen_kv, + is_causal=is_causal, + backend="cutlass-fmha", + ) + + @skip_if_libnatten_is_not_supported() + @skip_if_hopper_kernels_not_supported() + def test_cutlass_hopper_fmha_fast(self): + problem_sizes = [ + (3, 2, 64, 256, 7000), + (3, 2, 64, 256, 8192), + (3, 2, 64, 163, 11012), + (1, 1, 32, 32, 128), + (1, 1, 32, 128, 128), + (1, 1, 128, 128, 128), + (2, 1, 128, 128, 128), + (1, 2, 128, 128, 128), + (2, 2, 128, 128, 128), + (2, 2, 64, 128, 128), + (1, 1, 32, 32, 32), + (1, 1, 32, 128, 128), + (1, 1, 128, 3584, 381), + (1, 1, 128, 12072, 1680), + (1, 1, 32, 128, 128), + (1, 1, 32, 128, 4096), + (1, 1, 32, 128, 258), + (1, 2, 64, 128, 15), + (1, 1, 32, 8, 17), + (1, 1, 64, 17, 49), + (2, 4, 32, 128, 237), + (4, 3, 64, 256, 33), + (1, 1, 128, 128, 75), + (1, 1, 32, 125, 444), + (1, 2, 64, 125, 231), + (1, 1, 128, 256, 10240), + ] + for i, ( + batch, + heads, + head_dim, + seqlen_q, + seqlen_kv, + ) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for is_causal in [True, False]: + self._test_backend_against_torch_sdpa( + batch=batch, + heads=heads, + head_dim=head_dim, + seqlen_q=seqlen_q, + seqlen_kv=seqlen_kv, + is_causal=is_causal, + backend="hopper-fmha", + ) + self._test_backend_against_natten( + batch=batch, + heads=heads, + head_dim=head_dim, + seqlen_q=seqlen_q, + seqlen_kv=seqlen_kv, + is_causal=is_causal, + backend="hopper-fmha", + ) + + @skip_if_libnatten_is_not_supported() + @skip_if_blackwell_kernels_not_supported() + def test_cutlass_blackwell_fmha_fast(self): + problem_sizes = [ + # + (1, 1, 1, 128, 128, 128), + (1, 4, 1, 120, 128, 128), + (1, 4, 2, 72, 128, 128), + (1, 4, 4, 96, 128, 128), + (1, 1, 1, 48, 128, 1024), + (1, 1, 1, 128, 128, 13568), + (1, 1, 1, 16, 128, 13496), + (1, 1, 1, 32, 128, 13496), + (1, 1, 1, 32, 32, 13496), + (3, 1, 1, 32, 77, 8504), + (1, 1, 1, 32, 77, 8504), + (1, 1, 1, 64, 40, 12296), + (1, 2, 1, 64, 40, 12296), + (1, 2, 2, 64, 40, 12296), + (1, 1, 1, 128, 128, 128), + (2, 2, 1, 64, 128, 128), + (2, 2, 2, 64, 128, 128), + (1, 1, 1, 32, 32, 32), + (1, 1, 1, 32, 128, 128), + (1, 1, 1, 128, 3584, 381), + (1, 1, 1, 112, 12072, 1680), + (1, 1, 1, 32, 128, 128), + (1, 1, 1, 32, 128, 4096), + (1, 1, 1, 32, 128, 258), + (1, 2, 1, 64, 128, 15), + (1, 2, 2, 64, 128, 15), + (1, 1, 1, 32, 8, 17), + (1, 1, 1, 64, 17, 49), + (2, 4, 1, 32, 128, 237), + (2, 4, 2, 32, 128, 237), + (2, 4, 4, 32, 128, 237), + (4, 3, 1, 64, 256, 33), + (4, 3, 3, 64, 256, 33), + (1, 1, 1, 128, 128, 75), + (1, 1, 1, 32, 125, 444), + (1, 2, 1, 64, 125, 231), + (1, 2, 2, 64, 125, 231), + (1, 1, 1, 128, 256, 10240), + ] + for i, ( + batch, + heads_q, + heads_kv, + head_dim, + seqlen_q, + seqlen_kv, + ) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for is_causal in [False, True]: + self._test_backend_against_torch_sdpa( + batch=batch, + heads=heads_q, + heads_kv=heads_kv, + head_dim=head_dim, + seqlen_q=seqlen_q, + seqlen_kv=seqlen_kv, + is_causal=is_causal, + backend="blackwell-fmha", + ) + if heads_q == heads_kv: + self._test_backend_against_natten( + batch=batch, + heads=heads_q, + head_dim=head_dim, + seqlen_q=seqlen_q, + seqlen_kv=seqlen_kv, + is_causal=is_causal, + backend="blackwell-fmha", + ) + + @skip_if_libnatten_is_not_supported() + @skip_if_blackwell_kernels_not_supported() + def test_cutlass_blackwell_fmha_determinism(self): + problem_sizes = [ + (3, 7, 1, 128, 3584, 2077), + (2, 4, 2, 32, 1639, 4000), + (2, 4, 4, 32, 4096, 4096), + ] + + for i, ( + batch, + heads_q, + heads_kv, + head_dim, + seqlen_q, + seqlen_kv, + ) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for is_causal in [False, True]: + self._test_backend_against_natten( + batch=batch, + heads=heads_q, + heads_kv=heads_kv, + head_dim=head_dim, + seqlen_q=seqlen_q, + seqlen_kv=seqlen_kv, + is_causal=is_causal, + backend="blackwell-fmha", + determinism=True, + ) + + @pytest.mark.xfail( + reason="exact match must fail when determinism is disabled.", strict=True + ) + @skip_if_libnatten_is_not_supported() + @skip_if_blackwell_kernels_not_supported() + def test_cutlass_blackwell_fmha_determinism_xfail0(self): + _reset_everything(random_seed=0, torch_seed=0) + problem_sizes = [ + (3, 7, 1, 128, 3584, 2077), + ] + + for ( + batch, + heads_q, + heads_kv, + head_dim, + seqlen_q, + seqlen_kv, + ) in problem_sizes: + self._test_backend_against_natten( + batch=batch, + heads=heads_q, + heads_kv=heads_kv, + head_dim=head_dim, + seqlen_q=seqlen_q, + seqlen_kv=seqlen_kv, + is_causal=False, + backend="blackwell-fmha", + determinism=True, + fail_case=True, + ) + + @pytest.mark.xfail( + reason="exact match must fail when determinism is disabled.", strict=True + ) + @skip_if_libnatten_is_not_supported() + @skip_if_blackwell_kernels_not_supported() + def test_cutlass_blackwell_fmha_determinism_xfail1(self): + _reset_everything(random_seed=0, torch_seed=0) + problem_sizes = [ + (2, 4, 2, 32, 11283, 9873), + ] + + for ( + batch, + heads_q, + heads_kv, + head_dim, + seqlen_q, + seqlen_kv, + ) in problem_sizes: + self._test_backend_against_natten( + batch=batch, + heads=heads_q, + heads_kv=heads_kv, + head_dim=head_dim, + seqlen_q=seqlen_q, + seqlen_kv=seqlen_kv, + is_causal=True, + backend="blackwell-fmha", + determinism=True, + fail_case=True, + ) + + @skip_if_libnatten_is_not_supported() + def test_cutlass_fmha_determinism(self): + problem_sizes = [ + (3, 7, 1, 128, 3584, 2077), + (2, 4, 2, 32, 1639, 4000), + (2, 4, 4, 32, 4096, 4096), + ] + + for i, ( + batch, + heads_q, + heads_kv, + head_dim, + seqlen_q, + seqlen_kv, + ) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for is_causal in [False, True]: + self._test_backend_against_natten( + batch=batch, + heads=heads_q, + head_dim=head_dim, + seqlen_q=seqlen_q, + seqlen_kv=seqlen_kv, + is_causal=is_causal, + backend="cutlass-fmha", + determinism=True, + ) + + @skip_if_not_running_extended_tests() + @skip_if_libnatten_is_not_supported() + def test_cutlass_fmha_determinism_randsweep(self): + self._test_randsweep( + backend="cutlass-fmha", + max_tests=max(2, RAND_SWEEP_TESTS // 20), + reference="natten", + determinism=True, + ) + + @skip_if_not_running_extended_tests() + @skip_if_libnatten_is_not_supported() + def test_cutlass_fmha_randsweep_against_torch_sdpa(self): + self._test_randsweep( + backend="cutlass-fmha", max_tests=max(1, RAND_SWEEP_TESTS // 2) + ) + + @skip_if_not_running_extended_tests() + @skip_if_libnatten_is_not_supported() + @skip_if_blackwell_kernels_not_supported() + def test_cutlass_blackwell_fmha_randsweep_against_torch_sdpa(self): + self._test_randsweep(backend="blackwell-fmha", max_tests=RAND_SWEEP_TESTS) + + @skip_if_not_running_extended_tests() + @skip_if_libnatten_is_not_supported() + @skip_if_blackwell_kernels_not_supported() + def test_cutlass_blackwell_fmha_determinism_randsweep(self): + self._test_randsweep( + backend="blackwell-fmha", + max_tests=max(2, RAND_SWEEP_TESTS // 20), + reference="natten", + determinism=True, + ) + + @skip_if_not_running_extended_tests() + @skip_if_libnatten_is_not_supported() + @skip_if_hopper_kernels_not_supported() + def test_cutlass_hopper_fmha_randsweep_against_torch_sdpa(self): + self._test_randsweep(backend="hopper-fmha", max_tests=RAND_SWEEP_TESTS) + + +if __name__ == "__main__": + unittest.main() diff --git a/natten/tests/test_fmha_varlen.py b/natten/tests/test_fmha_varlen.py new file mode 100644 index 00000000..13f9c906 --- /dev/null +++ b/natten/tests/test_fmha_varlen.py @@ -0,0 +1,980 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +import random +import unittest +from typing import List, Optional, Tuple + +import natten # noqa: F401 +import torch +from natten._environment import _NUM_RAND_SWEEP_TESTS as RAND_SWEEP_TESTS +from natten.backends.configs.cutlass import ( + get_all_fmha_backward_configs, + get_all_fmha_forward_configs, +) +from natten.backends.configs.cutlass_blackwell import ( + get_all_fmha_backward_configs as get_all_blackwell_fmha_backward_configs, + get_all_fmha_forward_configs as get_all_blackwell_fmha_forward_configs, +) +from natten.backends.configs.cutlass_hopper import ( + get_all_fmha_backward_configs as get_all_hopper_fmha_backward_configs, + get_all_fmha_forward_configs as get_all_hopper_fmha_forward_configs, +) +from natten.functional import attention +from natten.types import KernelSchedule +from natten.utils.dtype import is_fp8 +from natten.utils.testing import ( + skip_if_blackwell_kernels_not_supported, + skip_if_hopper_kernels_not_supported, + skip_if_libnatten_is_not_supported, + skip_if_not_running_extended_tests, +) + +from .utils import logger + + +def _reset_everything(random_seed: int = 42, torch_seed: int = 42): + from natten.context import ( + NattenContext, + set_memory_usage_preference, + use_kv_parallelism_in_fused_na, + ) + + NattenContext.reset() + set_memory_usage_preference("unrestricted") + use_kv_parallelism_in_fused_na(True) + + random.seed(random_seed) + torch.manual_seed(torch_seed) + logger.debug(f"Reset seeds: {random_seed=}, {torch_seed=}") + torch.cuda.empty_cache() + torch.use_deterministic_algorithms(False) + + +def compute_split_reference( + batch: int, + heads: int, + head_dim: int, + seqlens_Q_list: List[int], + seqlens_KV_list: List[int], + is_causal: bool, + backend: str, + test_backprop: bool, + dtype: torch.dtype = torch.float32, + heads_kv: Optional[int] = None, + head_dim_v: Optional[int] = None, + q_tile_size: Optional[int] = None, + kv_tile_size: Optional[int] = None, + backward_q_tile_size: Optional[int] = None, + backward_kv_tile_size: Optional[int] = None, + backward_kv_splits: Optional[int] = None, + backward_use_pt_reduction: bool = False, + run_persistent_kernel: bool = True, + kernel_schedule: Optional[KernelSchedule] = None, +): + heads_kv = heads_kv or heads + head_dim_v = head_dim_v or head_dim + + assert len(seqlens_Q_list) == len(seqlens_KV_list) == batch + + seqlen_q_total = sum(seqlens_Q_list) + seqlen_kv_total = sum(seqlens_KV_list) + dtype_safe = torch.float16 + with torch.no_grad(): + q_ref, k_ref, v_ref, d_out_ref = ( + torch.randn( + (1, seqlen_q_total, heads, head_dim), device="cuda", dtype=dtype_safe + ).to(dtype), + torch.randn( + (1, seqlen_kv_total, heads_kv, head_dim), + device="cuda", + dtype=dtype_safe, + ).to(dtype), + torch.randn( + (1, seqlen_kv_total, heads_kv, head_dim_v), + device="cuda", + dtype=dtype_safe, + ).to(dtype), + torch.randn( + (1, seqlen_q_total, heads, head_dim_v), device="cuda", dtype=dtype_safe + ).to(dtype), + ) + q, k, v, d_out = ( + q_ref.clone(), + k_ref.clone(), + v_ref.clone(), + d_out_ref.clone(), + ) + + out_list = [] + lse_list = [] + d_q_list = [] + d_k_list = [] + d_v_list = [] + + q_start, kv_start = 0, 0 + for b in range(batch): + seqlen_q = seqlens_Q_list[b] + seqlen_kv = seqlens_KV_list[b] + + q_ = q_ref[:, q_start : q_start + seqlen_q, :, :].clone() + k_ = k_ref[:, kv_start : kv_start + seqlen_kv, :, :].clone() + v_ = v_ref[:, kv_start : kv_start + seqlen_kv, :, :].clone() + + if test_backprop: + q_ = q_.requires_grad_(True) + k_ = k_.requires_grad_(True) + v_ = v_.requires_grad_(True) + d_out_ = ( + d_out_ref[:, q_start : q_start + seqlen_q, :, :] + .clone() + .requires_grad_(True) + ) + + out_, lse_ = attention( + q_, + k_, + v_, + is_causal=is_causal, + backend=backend, + return_lse=True, + q_tile_size=q_tile_size, + kv_tile_size=kv_tile_size, + backward_q_tile_size=backward_q_tile_size, + backward_kv_tile_size=backward_kv_tile_size, + backward_kv_splits=backward_kv_splits, + backward_use_pt_reduction=backward_use_pt_reduction, + run_persistent_kernel=run_persistent_kernel, + kernel_schedule=kernel_schedule, + ) + + if test_backprop: + out_.backward(d_out_) + + with torch.no_grad(): + out_list.append(out_.data.clone().float()) + lse_list.append(lse_.data.clone().float()) + if test_backprop: + assert q_.grad is not None + assert k_.grad is not None + assert v_.grad is not None + d_q_list.append(q_.grad.clone().float()) + d_k_list.append(k_.grad.clone().float()) + d_v_list.append(v_.grad.clone().float()) + + q_start += seqlen_q + kv_start += seqlen_kv + + assert q_start == seqlen_q_total + assert kv_start == seqlen_kv_total + + out_ref = torch.cat(out_list, dim=1) + lse_ref = torch.cat(lse_list, dim=1) + assert out_ref.shape[:3] == q_ref.shape[:3] + dq_ref = None + dk_ref = None + dv_ref = None + if test_backprop: + dq_ref = torch.cat(d_q_list, dim=1) + dk_ref = torch.cat(d_k_list, dim=1) + dv_ref = torch.cat(d_v_list, dim=1) + + assert dq_ref.shape == q_ref.shape + assert dk_ref.shape == k_ref.shape + assert dv_ref.shape == v_ref.shape + + return (q, k, v, d_out), (out_ref, lse_ref, dq_ref, dk_ref, dv_ref) + + +class FMHAVarlenTest(unittest.TestCase): + def setUp(self): + _reset_everything() + + def tearDown(self): + _reset_everything() + + def _test_against_manual_varlen( + self, + batch: int, + heads: int, + head_dim: int, + seqlens_Q_list: List[int], + seqlens_KV_list: List[int], + is_causal: bool, + dtype: torch.dtype, + atol_fwd: Tuple[float, float], + atol_bwd: Optional[Tuple[float, float, float]], + backend: str, + reference_backend: str, + test_backprop: bool, + q_tile_size: Optional[int] = None, + kv_tile_size: Optional[int] = None, + backward_q_tile_size: Optional[int] = None, + backward_kv_tile_size: Optional[int] = None, + backward_kv_splits: Optional[int] = None, + backward_use_pt_reduction: bool = False, + run_persistent_kernel: bool = True, + kernel_schedule: Optional[KernelSchedule] = None, + reference_q_tile_size: Optional[int] = None, + reference_kv_tile_size: Optional[int] = None, + reference_backward_q_tile_size: Optional[int] = None, + reference_backward_kv_tile_size: Optional[int] = None, + reference_backward_kv_splits: Optional[int] = None, + reference_backward_use_pt_reduction: bool = False, + reference_run_persistent_kernel: bool = True, + heads_kv: Optional[int] = None, + head_dim_v: Optional[int] = None, + ): + heads_kv = heads_kv or heads + head_dim_v = head_dim_v or head_dim + + determinism = torch.are_deterministic_algorithms_enabled() + logger.debug( + f"Testing FMHA varlen ({backend}) vs {reference_backend}: {batch=}, {heads=}, {heads_kv=}, {head_dim=}, {head_dim_v=}, " + f"{seqlens_Q_list=}, {seqlens_KV_list=}, {is_causal=}, {dtype=}, " + f"{q_tile_size=}, {kv_tile_size=}, {run_persistent_kernel=}, {kernel_schedule=}" + + ( + f",\n({determinism=}) {backward_q_tile_size=}, {backward_kv_tile_size=}, " + f"{backward_kv_splits=}, {backward_use_pt_reduction=}." + if test_backprop + else "." + ) + ) + + inputs, reference = compute_split_reference( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + head_dim_v=head_dim_v, + seqlens_Q_list=seqlens_Q_list, + seqlens_KV_list=seqlens_KV_list, + is_causal=is_causal, + dtype=dtype, + backend=reference_backend, + q_tile_size=reference_q_tile_size, + kv_tile_size=reference_kv_tile_size, + backward_q_tile_size=reference_backward_q_tile_size, + backward_kv_tile_size=reference_backward_kv_tile_size, + backward_kv_splits=reference_backward_kv_splits, + backward_use_pt_reduction=reference_backward_use_pt_reduction, + run_persistent_kernel=reference_run_persistent_kernel, + kernel_schedule=kernel_schedule, + test_backprop=test_backprop, + ) + + q, k, v, d_out = inputs + out_ref, lse_ref, dq_ref, dk_ref, dv_ref = reference + q = q.to(dtype) + k = k.to(dtype) + v = v.to(dtype) + d_out = d_out.to(dtype) + + # Run target + if test_backprop: + q.requires_grad_(test_backprop) + k.requires_grad_(test_backprop) + v.requires_grad_(test_backprop) + d_out.requires_grad_(test_backprop) + + seqlens_Q = torch.tensor(seqlens_Q_list, dtype=torch.int32, device=q.device) + seqlens_KV = torch.tensor(seqlens_KV_list, dtype=torch.int32, device=q.device) + + out_, lse_ = attention( + q, + k, + v, + is_causal=is_causal, + backend=backend, + q_tile_size=q_tile_size, + kv_tile_size=kv_tile_size, + backward_q_tile_size=backward_q_tile_size, + backward_kv_tile_size=backward_kv_tile_size, + backward_kv_splits=backward_kv_splits, + backward_use_pt_reduction=backward_use_pt_reduction, + run_persistent_kernel=run_persistent_kernel, + kernel_schedule=kernel_schedule, + return_lse=True, + seqlens_Q=seqlens_Q, + seqlens_KV=seqlens_KV, + ) + out = out_.float() + lse = lse_.float() + + if test_backprop: + dq, dk, dv = None, None, None + out_.backward(d_out) + with torch.no_grad(): + dq, dk, dv = ( + q.grad.clone().float(), + k.grad.clone().float(), + v.grad.clone().float(), + ) + + atol_out, atol_lse = atol_fwd + assert out.shape == out_ref.shape + + torch.testing.assert_close(out, out_ref, atol=atol_out, rtol=0) + torch.testing.assert_close(lse, lse_ref, atol=atol_lse, rtol=0) + + if test_backprop: + assert atol_bwd is not None + atol_dq, atol_dk, atol_dv = atol_bwd + torch.testing.assert_close(dq, dq_ref, atol=atol_dq, rtol=0) + torch.testing.assert_close(dk, dk_ref, atol=atol_dk, rtol=0) + torch.testing.assert_close(dv, dv_ref, atol=atol_dv, rtol=0) + + def _test_cutlass_fmha_varlen( + self, + batch, + heads, + head_dim, + head_dim_v, + seqlens_Q_list, + seqlens_KV_list, + is_causal, + max_configs=5, + determinism=False, + fail_case=False, + ): + head_dim_v = head_dim_v or head_dim + torch.set_default_device("cuda") + + # We're testing against the same backend and same dtype, but implemented with for loops over + # different sequences. Given how the typical varlen scheduler operates, we expect exact + # matches over everything except dQ, which is reduced. + # This is therefore only a test of the varlen functionality. + # Correctness per dtype is expected to be verified in the main fmha tests. + if determinism: + torch.use_deterministic_algorithms(not fail_case) + ALLOWED_DTYPES = [ + # dtype, (atol_out, atol_lse), (atol_dq, atol_dk, atol_dv) + (torch.float32, (0, 0), (0, 0, 0)), + (torch.float16, (0, 0), (0, 0, 0)), + (torch.bfloat16, (0, 0), (0, 0, 0)), + ] + else: + ALLOWED_DTYPES = [ + # dtype, (atol_out, atol_lse), (atol_dq, atol_dk, atol_dv) + (torch.float32, (0, 0), (1e-2, 0, 0)), + (torch.float16, (0, 0), (1e-2, 0, 0)), + (torch.bfloat16, (0, 0), (1e-2, 0, 0)), + ] + + for dtype, atol_fwd, atol_bwd in ALLOWED_DTYPES: + dummy_fwd = torch.empty( + (1, min(seqlens_Q_list), heads, max(head_dim, head_dim_v)), + device="cuda", + dtype=dtype, + ) + dummy_bwd = torch.empty( + (1, min(seqlens_KV_list), heads, max(head_dim, head_dim_v)), + device="cuda", + dtype=dtype, + ) + + forward_configs = get_all_fmha_forward_configs(dummy_fwd) + backward_configs = get_all_fmha_backward_configs(dummy_bwd) + assert len(forward_configs) > 0 + assert len(backward_configs) > 0 + + random.shuffle(forward_configs) + random.shuffle(backward_configs) + + n_configs_to_test = min( + max_configs, max(len(forward_configs), len(backward_configs)) + ) + + for i in range(n_configs_to_test): + q_tile_size, kv_tile_size = forward_configs[i % len(forward_configs)] + backward_q_tile_size, backward_kv_tile_size = backward_configs[ + i % len(backward_configs) + ] + + seqlen_kv = dummy_bwd.shape[1] + kv_splits = ( + seqlen_kv + backward_kv_tile_size - 1 + ) // backward_kv_tile_size + + self._test_against_manual_varlen( + batch=batch, + heads=heads, + head_dim=head_dim, + head_dim_v=head_dim_v, + seqlens_Q_list=seqlens_Q_list, + seqlens_KV_list=seqlens_KV_list, + is_causal=is_causal, + dtype=dtype, + atol_fwd=atol_fwd, + atol_bwd=atol_bwd, + backend="cutlass-fmha", + reference_backend="cutlass-fmha", + test_backprop=True, + q_tile_size=q_tile_size, + kv_tile_size=kv_tile_size, + backward_q_tile_size=backward_q_tile_size, + backward_kv_tile_size=backward_kv_tile_size, + backward_kv_splits=kv_splits, + backward_use_pt_reduction=False, + reference_q_tile_size=q_tile_size, + reference_kv_tile_size=kv_tile_size, + reference_backward_q_tile_size=backward_q_tile_size, + reference_backward_kv_tile_size=backward_kv_tile_size, + reference_backward_kv_splits=kv_splits, + reference_backward_use_pt_reduction=False, + ) + + torch.use_deterministic_algorithms(False) + + def _test_cutlass_hopper_fmha_varlen( + self, + batch, + heads, + head_dim, + seqlens_Q_list, + seqlens_KV_list, + is_causal, + heads_kv: Optional[int] = None, + ): + torch.set_default_device("cuda") + + # We're testing against the same backend and same dtype, but implemented with for loops over + # different sequences. Given how the typical varlen scheduler operates, we expect exact + # matches over everything except dQ, which is reduced. + # This is therefore only a test of the varlen functionality. + # Correctness per dtype is expected to be verified in the main + # fmha tests. + ALLOWED_DTYPES = [ + # dtype, (atol_out, atol_lse), (atol_dq, atol_dk, atol_dv) + (torch.float16, (0, 0), (1e-2, 0, 0)), + (torch.bfloat16, (0, 0), (1e-2, 0, 0)), + ] + + for dtype, atol_fwd, atol_bwd in ALLOWED_DTYPES: + dummy = torch.empty((1, 128, heads, head_dim), device="cuda", dtype=dtype) + + forward_configs = get_all_hopper_fmha_forward_configs(dummy) + backward_configs = get_all_hopper_fmha_backward_configs(dummy) + assert len(forward_configs) > 0 + assert len(backward_configs) > 0 + + random.shuffle(forward_configs) + random.shuffle(backward_configs) + + for i in range(max(len(forward_configs), len(backward_configs))): + (q_tile_size, kv_tile_size), kernel_schedule = forward_configs[ + i % len(forward_configs) + ] + backward_q_tile_size, backward_kv_tile_size = backward_configs[ + i % len(backward_configs) + ] + + self._test_against_manual_varlen( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + seqlens_Q_list=seqlens_Q_list, + seqlens_KV_list=seqlens_KV_list, + is_causal=is_causal, + dtype=dtype, + atol_fwd=atol_fwd, + atol_bwd=atol_bwd, + backend="hopper-fmha", + reference_backend="hopper-fmha", + test_backprop=True, + q_tile_size=q_tile_size, + kv_tile_size=kv_tile_size, + backward_q_tile_size=backward_q_tile_size, + backward_kv_tile_size=backward_kv_tile_size, + kernel_schedule=kernel_schedule, + ) + + def _test_cutlass_blackwell_fmha_varlen( + self, + batch, + heads, + head_dim, + seqlens_Q_list, + seqlens_KV_list, + is_causal, + heads_kv: Optional[int] = None, + determinism=False, + fail_case=False, + ): + torch.set_default_device("cuda") + + # We're testing against the same backend and same dtype, but implemented with for loops over + # different sequences. Given how the typical varlen scheduler operates, we expect exact + # matches over everything except dQ, which is reduced. + # This is therefore only a test of the varlen functionality. + # Correctness per dtype is expected to be verified in the main fmha tests. + ALLOWED_DTYPES = [ + # dtype, (atol_out, atol_lse), (atol_dq, atol_dk, atol_dv) + (torch.float16, (0, 0), (1e-2, 0, 0)), + (torch.bfloat16, (0, 0), (1e-2, 0, 0)), + ] + + if determinism: + torch.use_deterministic_algorithms(not fail_case) + ALLOWED_DTYPES = [ + # dtype, (atol_out, atol_lse), (atol_dq, atol_dk, atol_dv) + (torch.float16, (0, 0), (0, 0, 0)), + (torch.bfloat16, (0, 0), (0, 0, 0)), + ] + else: + ALLOWED_DTYPES += [ + (torch.float8_e4m3fn, (0, 0), None), # type: ignore[list-item] + (torch.float8_e5m2, (0, 0), None), # type: ignore[list-item] + ] + + for dtype, atol_fwd, atol_bwd in ALLOWED_DTYPES: + if is_fp8(dtype) and head_dim % 16 != 0: + continue + + dummy = torch.empty((1, 128, heads, head_dim), device="cuda", dtype=dtype) + + forward_configs = get_all_blackwell_fmha_forward_configs(dummy) + backward_configs = get_all_blackwell_fmha_backward_configs(dummy) + assert len(forward_configs) > 0 + assert len(backward_configs) > 0 + + random.shuffle(forward_configs) + random.shuffle(backward_configs) + + for i in range(max(len(forward_configs), len(backward_configs))): + q_tile_size, kv_tile_size = forward_configs[i % len(forward_configs)] + backward_q_tile_size, backward_kv_tile_size = backward_configs[ + i % len(backward_configs) + ] + + for run_persistent_kernel in [True, False]: + self._test_against_manual_varlen( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + seqlens_Q_list=seqlens_Q_list, + seqlens_KV_list=seqlens_KV_list, + is_causal=is_causal, + dtype=dtype, + atol_fwd=atol_fwd, + atol_bwd=atol_bwd, + backend="blackwell-fmha", + reference_backend="blackwell-fmha", + test_backprop=not is_fp8(dtype), + q_tile_size=q_tile_size, + kv_tile_size=kv_tile_size, + backward_q_tile_size=backward_q_tile_size, + backward_kv_tile_size=backward_kv_tile_size, + run_persistent_kernel=run_persistent_kernel, + ) + + torch.use_deterministic_algorithms(False) + + def _test_varlen_backend( + self, + batch, + heads, + head_dim, + seqlens_Q_list, + seqlens_KV_list, + is_causal, + backend, + heads_kv=None, + head_dim_v=None, + determinism: bool = False, + fail_case: bool = False, + ): + if backend == "cutlass-fmha": + assert heads_kv is None or heads_kv == heads + return self._test_cutlass_fmha_varlen( + batch=batch, + heads=heads, + head_dim=head_dim, + head_dim_v=head_dim_v, + seqlens_Q_list=seqlens_Q_list, + seqlens_KV_list=seqlens_KV_list, + is_causal=is_causal, + determinism=determinism, + fail_case=fail_case, + ) + elif backend == "hopper-fmha": + assert not determinism + assert not fail_case + assert ( + head_dim_v is None or head_dim_v == head_dim + ), "Hopper FMHA does not allow head_dim_v." + assert heads_kv is None or heads_kv == heads + + return self._test_cutlass_hopper_fmha_varlen( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + seqlens_Q_list=seqlens_Q_list, + seqlens_KV_list=seqlens_KV_list, + is_causal=is_causal, + ) + elif backend == "blackwell-fmha": + assert ( + head_dim_v is None or head_dim_v == head_dim + ), "Blackwell FMHA does not allow head_dim_v." + + return self._test_cutlass_blackwell_fmha_varlen( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + seqlens_Q_list=seqlens_Q_list, + seqlens_KV_list=seqlens_KV_list, + is_causal=is_causal, + determinism=determinism, + fail_case=fail_case, + ) + else: + raise NotImplementedError(f"Add {backend=} to tests.") + + def _test_varlen_randsweep( + self, backend, max_tests=1000, determinism: bool = False + ): + max_qk = 2**17 + for i in range(max_tests): + # to help with reproducibility of use cases + _reset_everything(random_seed=i, torch_seed=i) + batch = random.choice(range(1, 10)) + + supports_dim_v = False + supports_gqa_mqa = False + if backend == "blackwell-fmha": + head_dim_choices = list(range(8, 129, 8)) + heads_choices = range(1, 8 + 1) + supports_gqa_mqa = True + elif backend == "hopper-fmha": + head_dim_choices = [32, 64, 128] + heads_choices = range(1, 4) + else: + assert backend == "cutlass-fmha" + head_dim_choices = range(8, 256 + 1, 8) # type: ignore[assignment] + heads_choices = range(1, 4) + supports_dim_v = True + + heads = random.choice(heads_choices) + heads_kv = ( + heads + if not supports_gqa_mqa + else random.choice( + [1] + [i for i in range(1, heads + 1) if heads % i == 0] + ) + ) + assert heads >= heads_kv and heads % heads_kv == 0 + + head_dim = random.choice(head_dim_choices) + head_dim_v = ( + head_dim if not supports_dim_v else random.choice(head_dim_choices) + ) + + seqlens_Q_list = [] # type: ignore[var-annotated] + seqlens_KV_list = [] # type: ignore[var-annotated] + for i in range(batch): + max_q = min(2**12, max(max_qk - sum(seqlens_Q_list), 24)) + max_k = min(2**12, max(max_qk - sum(seqlens_KV_list), 24)) + new_q = random.choice(range(1, max_q, 1)) + new_k = random.choice(range(8, max_k, 1)) + seqlens_Q_list.append(new_q) + seqlens_KV_list.append(new_k) + + for is_causal in [False, True]: + self._test_varlen_backend( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + head_dim_v=head_dim_v, + seqlens_Q_list=seqlens_Q_list, + seqlens_KV_list=seqlens_KV_list, + is_causal=is_causal, + backend=backend, + determinism=determinism, + ) + + @skip_if_libnatten_is_not_supported() + def test_cutlass_varlen_fmha_fast(self): + problem_sizes = [ + (4, 1, 88, [11, 1846, 2433, 1193], [1681, 2831, 1932, 141]), + ( + 7, + 7, + 56, + [126, 1127, 3130, 3782, 1445, 751, 594], + [746, 2886, 1273, 1399, 33, 3570, 2327], + ), + ( + 8, + 6, + 240, + [1912, 1057, 1518, 2489, 3955, 2950, 3193, 678], + [2873, 3149, 695, 2495, 3071, 477, 3363, 3959], + ), + (6, 1, 128, [128, 128, 135, 121, 128, 128], [128, 128, 135, 121, 128, 128]), + (5, 1, 128, [128, 128, 135, 128, 128], [128, 128, 135, 128, 128]), + (2, 1, 128, [135, 200], [128, 768]), + (2, 1, 128, [1024, 200], [128, 768]), + (2, 1, 128, [135, 200], [135, 768]), + (2, 1, 128, [1024, 200], [135, 768]), + (2, 1, 128, [1024, 256], [128, 768]), + (4, 1, 128, [1024, 8, 17, 2048], [10, 20, 512, 16]), + (3, 2, 128, [268, 1584, 1571], [2448, 4088, 1925]), + (2, 1, 128, [1024, 256], [512, 768]), + ] + for i, ( + batch, + heads, + head_dim, + seqlens_Q_list, + seqlens_KV_list, + ) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for is_causal in [False, True]: + self._test_varlen_backend( + batch=batch, + heads=heads, + head_dim=head_dim, + seqlens_Q_list=seqlens_Q_list, + seqlens_KV_list=seqlens_KV_list, + is_causal=is_causal, + backend="cutlass-fmha", + ) + + @skip_if_not_running_extended_tests() + @skip_if_libnatten_is_not_supported() + def test_cutlass_varlen_fmha_extended(self): + self._test_varlen_randsweep(backend="cutlass-fmha", max_tests=RAND_SWEEP_TESTS) + + @skip_if_libnatten_is_not_supported() + def test_cutlass_varlen_determinism_fmha_fast(self): + problem_sizes = [ + ( + 9, + 4, + 128, + [2669, 2240, 910, 2421, 3323, 34, 3308, 2867, 1401], + [2880, 1726, 1847, 1147, 3568, 3116, 661, 1739, 1146], + ), + (6, 1, 128, [128, 128, 135, 121, 128, 128], [128, 128, 135, 121, 128, 128]), + (5, 1, 128, [128, 128, 135, 128, 128], [128, 128, 135, 128, 128]), + (2, 1, 128, [135, 200], [128, 768]), + (2, 1, 128, [1024, 200], [128, 768]), + (2, 1, 128, [135, 200], [135, 768]), + (2, 1, 128, [1024, 200], [135, 768]), + (2, 1, 128, [1024, 256], [128, 768]), + (4, 1, 128, [1024, 8, 17, 2048], [10, 20, 512, 16]), + (3, 2, 128, [268, 1584, 1571], [2448, 4088, 1925]), + (2, 1, 128, [1024, 256], [512, 768]), + ] + for i, ( + batch, + heads, + head_dim, + seqlens_Q_list, + seqlens_KV_list, + ) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for is_causal in [False, True]: + self._test_varlen_backend( + batch=batch, + heads=heads, + head_dim=head_dim, + seqlens_Q_list=seqlens_Q_list, + seqlens_KV_list=seqlens_KV_list, + is_causal=is_causal, + backend="cutlass-fmha", + determinism=True, + ) + + @skip_if_not_running_extended_tests() + @skip_if_libnatten_is_not_supported() + def test_cutlass_varlen_fmha_determinism_extended(self): + self._test_varlen_randsweep( + backend="cutlass-fmha", + max_tests=max(2, RAND_SWEEP_TESTS // 100), + determinism=True, + ) + + @skip_if_libnatten_is_not_supported() + @skip_if_hopper_kernels_not_supported() + def test_cutlass_hopper_varlen_fmha_fast(self): + problem_sizes = [ + ( + 9, + 4, + 128, + [2669, 2240, 910, 2421, 3323, 34, 3308, 2867, 1401], + [2880, 1726, 1847, 1147, 3568, 3116, 661, 1739, 1146], + ), + (6, 1, 128, [128, 128, 135, 121, 128, 128], [128, 128, 135, 121, 128, 128]), + (5, 1, 128, [128, 128, 135, 128, 128], [128, 128, 135, 128, 128]), + (2, 1, 128, [135, 200], [128, 768]), + (2, 1, 128, [1024, 200], [128, 768]), + (2, 1, 128, [135, 200], [135, 768]), + (2, 1, 128, [1024, 200], [135, 768]), + (2, 1, 128, [1024, 256], [128, 768]), + (4, 1, 128, [1024, 8, 17, 2048], [10, 20, 512, 16]), + (3, 2, 128, [268, 1584, 1571], [2448, 4088, 1925]), + (2, 1, 128, [16, 16], [16, 16]), + (2, 1, 128, [47, 17], [37, 32]), + (2, 1, 128, [48, 16], [37, 32]), + (2, 1, 64, [48, 16], [37, 32]), + (2, 1, 64, [64, 64], [64, 64]), + (2, 1, 64, [128, 128], [128, 128]), + (2, 1, 128, [512, 512], [512, 512]), + ] + for i, ( + batch, + heads, + head_dim, + seqlens_Q_list, + seqlens_KV_list, + ) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for is_causal in [False, True]: + self._test_varlen_backend( + batch=batch, + heads=heads, + head_dim=head_dim, + seqlens_Q_list=seqlens_Q_list, + seqlens_KV_list=seqlens_KV_list, + is_causal=is_causal, + backend="hopper-fmha", + ) + + @skip_if_not_running_extended_tests() + @skip_if_libnatten_is_not_supported() + @skip_if_hopper_kernels_not_supported() + def test_cutlass_hopper_varlen_fmha_extended(self): + self._test_varlen_randsweep(backend="hopper-fmha", max_tests=RAND_SWEEP_TESTS) + + @skip_if_libnatten_is_not_supported() + @skip_if_blackwell_kernels_not_supported() + def test_cutlass_blackwell_varlen_determinism_fmha_fast(self): + problem_sizes = [ + ( + 9, + 4, + 128, + [2669, 2240, 910, 2421, 3323, 34, 3308, 2867, 1401], + [2880, 1726, 1847, 1147, 3568, 3116, 661, 1739, 1146], + ), + (6, 1, 128, [128, 128, 135, 121, 128, 128], [128, 128, 135, 121, 128, 128]), + (5, 1, 128, [128, 128, 135, 128, 128], [128, 128, 135, 128, 128]), + (2, 1, 128, [135, 200], [128, 768]), + (2, 1, 128, [1024, 200], [128, 768]), + (2, 1, 128, [135, 200], [135, 768]), + (2, 1, 128, [1024, 200], [135, 768]), + (2, 1, 128, [1024, 256], [128, 768]), + (4, 1, 128, [1024, 8, 17, 2048], [10, 20, 512, 16]), + (3, 2, 128, [268, 1584, 1571], [2448, 4088, 1925]), + (2, 1, 128, [1024, 256], [512, 768]), + ] + for i, ( + batch, + heads, + head_dim, + seqlens_Q_list, + seqlens_KV_list, + ) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for is_causal in [False, True]: + self._test_varlen_backend( + batch=batch, + heads=heads, + head_dim=head_dim, + seqlens_Q_list=seqlens_Q_list, + seqlens_KV_list=seqlens_KV_list, + is_causal=is_causal, + backend="blackwell-fmha", + determinism=True, + ) + + @skip_if_libnatten_is_not_supported() + @skip_if_blackwell_kernels_not_supported() + def test_cutlass_blackwell_varlen_fmha_fast(self): + problem_sizes = [ + ( + 9, + 4, + 128, + [2669, 2240, 910, 2421, 3323, 34, 3308, 2867, 1401], + [2880, 1726, 1847, 1147, 3568, 3116, 661, 1739, 1146], + ), + (6, 1, 128, [128, 128, 135, 121, 128, 128], [128, 128, 135, 121, 128, 128]), + (5, 1, 128, [128, 128, 135, 128, 128], [128, 128, 135, 128, 128]), + (2, 1, 128, [135, 200], [128, 768]), + (2, 1, 128, [1024, 200], [128, 768]), + (2, 1, 128, [135, 200], [135, 768]), + (2, 1, 128, [1024, 200], [135, 768]), + (2, 1, 128, [1024, 256], [128, 768]), + (4, 1, 128, [1024, 8, 17, 2048], [10, 20, 512, 16]), + (3, 2, 128, [268, 1584, 1571], [2448, 4088, 1925]), + (2, 1, 128, [1024, 256], [512, 768]), + ] + for i, ( + batch, + heads, + head_dim, + seqlens_Q_list, + seqlens_KV_list, + ) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for is_causal in [False, True]: + self._test_varlen_backend( + batch=batch, + heads=heads, + head_dim=head_dim, + seqlens_Q_list=seqlens_Q_list, + seqlens_KV_list=seqlens_KV_list, + is_causal=is_causal, + backend="blackwell-fmha", + ) + + @skip_if_not_running_extended_tests() + @skip_if_libnatten_is_not_supported() + @skip_if_blackwell_kernels_not_supported() + def test_cutlass_blackwell_varlen_fmha_extended(self): + self._test_varlen_randsweep( + backend="blackwell-fmha", max_tests=RAND_SWEEP_TESTS + ) + + @skip_if_not_running_extended_tests() + @skip_if_libnatten_is_not_supported() + @skip_if_blackwell_kernels_not_supported() + def test_cutlass_blackwell_varlen_fmha_determinism_extended(self): + self._test_varlen_randsweep( + backend="blackwell-fmha", + max_tests=max(2, RAND_SWEEP_TESTS // 100), + determinism=True, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/natten/tests/test_fna.py b/natten/tests/test_fna.py new file mode 100644 index 00000000..9fe358f4 --- /dev/null +++ b/natten/tests/test_fna.py @@ -0,0 +1,791 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +import math +import random +import unittest +from itertools import product + +import natten # noqa: F401 +import pytest +import torch +from natten._environment import _NUM_RAND_SWEEP_TESTS as RAND_SWEEP_TESTS +from natten.backends.configs.cutlass import ( + get_all_backward_configs, + get_all_forward_configs, +) +from natten.utils.testing import ( + skip_if_libnatten_is_not_supported, + skip_if_not_running_extended_tests, + supports_bfloat16, + supports_float16, +) + +from .utils import logger, NattenBackendTester + + +def _reset_everything(random_seed: int = 42, torch_seed: int = 42): + from natten.context import ( + NattenContext, + set_memory_usage_preference, + use_kv_parallelism_in_fused_na, + ) + + NattenContext.reset() + set_memory_usage_preference("unrestricted") + use_kv_parallelism_in_fused_na(True) + + random.seed(random_seed) + torch.manual_seed(torch_seed) + logger.debug(f"Reset seeds: {random_seed=}, {torch_seed=}") + torch.cuda.empty_cache() + torch.use_deterministic_algorithms(False) + + +class FNABackendTest(unittest.TestCase): + def setUp(self): + _reset_everything() + + def tearDown(self): + _reset_everything() + + def _test_all_dtypes_against_reference( + self, + batch, + heads, + head_dim, + head_dim_v, + input_shape, + kernel_size, + stride, + dilation, + is_causal=None, + configs_to_test=5, + heads_kv=None, + ): + + torch.set_default_device("cuda") + assert isinstance(input_shape, tuple) + na_dim = len(input_shape) + assert na_dim in [1, 2, 3], "Only supports NA1D, 2D, 3D." + + tester = NattenBackendTester( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + head_dim_v=head_dim_v, + input_shape=input_shape, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + test_backprop=True, + reference_backend="reference", + reference_fmha_backend="reference", + dtype=torch.float32, + ) + + ALLOWED_DTYPES = [ + ( + torch.float32, + (1e-4, (1e-2, 1e-4, 1e-4)), + ), + ] + + if supports_float16(torch.get_default_device()): + ALLOWED_DTYPES.append( + ( + torch.float16, + (1e-2, (1e-2, 1e-2, 1e-2)), + ) + ) + + if supports_bfloat16(torch.get_default_device()): + ALLOWED_DTYPES.append( + ( + torch.bfloat16, + (5e-2, (1e-2, 1e-2, 1e-2)), + ) + ) + + for dtype, eps in ALLOWED_DTYPES: + + dummy = torch.randn( + (batch, *input_shape, heads, max(head_dim, head_dim_v)), + device="cuda", + dtype=dtype, + ) + + forward_configs = get_all_forward_configs(dummy) + backward_configs = get_all_backward_configs(dummy) + assert len(forward_configs) > 0 + assert len(backward_configs) > 0 + + num_fwd_configs = len(forward_configs) + num_bwd_configs = len(backward_configs) + + if num_fwd_configs * num_bwd_configs < configs_to_test: + cfg_idxes = product( + [x for x in range(num_fwd_configs)], + [x for x in range(num_bwd_configs)], + ) + else: + cfg_idxes = [ + (x % num_fwd_configs, x % num_bwd_configs) + for x in range(configs_to_test) + ] + + random.shuffle(forward_configs) + random.shuffle(backward_configs) + + for i, j in cfg_idxes: + q_tile_shape, kv_tile_shape = forward_configs[i] + backward_q_tile_shape, backward_kv_tile_shape = backward_configs[j] + + tester.test( + eps=eps, + dtype=dtype, + target_backend="cutlass-fna", + target_fmha_backend="cutlass-fmha", + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + backward_q_tile_shape=backward_q_tile_shape, + backward_kv_tile_shape=backward_kv_tile_shape, + ) + + @skip_if_libnatten_is_not_supported() + def test_1d_against_reference(self): + problem_sizes = [ + (4, 3, 3, 128, (256,), (255,), (82,), (1)), + (4, 3, 1, 128, (256,), (255,), (82,), (1)), + (2, 4, 4, 128, (128,), (63,), (31,), (1)), + (2, 4, 2, 128, (128,), (63,), (31,), (1)), + (2, 4, 1, 128, (128,), (63,), (31,), (1)), + (1, 1, 1, 128, (2048,), (2048,), (1,), (1)), + (1, 1, 1, 128, (32768,), (2048,), (2048,), (1)), + (1, 1, 1, 128, (32768,), (2048,), (256,), (1)), + (1, 1, 1, 128, (32768,), (2048,), (128,), (1)), + (1, 1, 1, 128, (32768,), (2048,), (1,), (1)), + (1, 1, 1, 32, (128,), (3,), (1,), (1)), + (1, 1, 1, 32, (128,), (3,), (2,), (5)), + (1, 2, 2, 64, (128,), (15,), (1,), (1)), + (1, 1, 1, 32, (128,), (3,), (2,), (10)), + (1, 1, 1, 64, (128,), (8,), (7,), (5)), + (1, 1, 1, 128, (128,), (61,), (33,), (1)), + (1, 1, 1, 32, (125,), (3,), (1,), (1)), + (1, 2, 2, 64, (125,), (15,), (1,), (1)), + (1, 1, 1, 128, (256,), (3,), (2,), (10)), + ] + for i, ( + batch, + heads, + heads_kv, + head_dim, + input_shape, + kernel_size, + stride, + dilation, + ) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for head_dim_v in [random.choice(range(8, 193, 8)), head_dim]: + for causal in [False, True]: + is_causal = (causal,) + self._test_all_dtypes_against_reference( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + head_dim_v=head_dim_v, + input_shape=input_shape, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + ) + + @skip_if_libnatten_is_not_supported() + def test_2d_against_reference(self): + problem_sizes = [ + (1, 4, 4, 32, (5, 25), (5, 16), (1, 16), (1, 1)), + (1, 4, 2, 32, (5, 25), (5, 16), (1, 16), (1, 1)), + (1, 4, 1, 32, (5, 25), (5, 16), (1, 16), (1, 1)), + (1, 1, 1, 32, (84, 19), (7, 3), (1, 1), (5, 1)), + (1, 1, 1, 128, (19, 29), (8, 8), (1, 1), (2, 3)), + (1, 1, 1, 128, (48, 17), (24, 16), (1, 1), (2, 1)), + (1, 1, 1, 128, (67, 80), (12, 7), (8, 4), (5, 11)), + (1, 1, 1, 128, (8, 8), (8, 8), (1, 1), (1, 1)), + (1, 1, 1, 128, (33, 33), (24, 16), (1, 1), (1, 1)), + (1, 1, 1, 32, (16, 16), (16, 16), (1, 1), (1, 1)), + (1, 1, 1, 128, (44, 80), (44, 80), (1, 1), (1, 1)), + (1, 1, 1, 32, (40, 20), (3, 7), (1, 1), (1, 1)), + (1, 1, 1, 32, (16, 16), (3, 3), (1, 1), (1, 1)), + (1, 1, 1, 128, (44, 80), (9, 10), (1, 1), (1, 1)), + (1, 1, 1, 64, (28, 40), (17, 31), (1, 1), (1, 1)), + (1, 1, 1, 64, (36, 40), (36, 40), (12, 13), (1, 1)), + (1, 1, 1, 128, (44, 80), (44, 80), (4, 8), (1, 1)), + ] + for i, ( + batch, + heads, + heads_kv, + head_dim, + input_shape, + kernel_size, + stride, + dilation, + ) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for head_dim_v in [random.choice(range(8, 193, 8)), head_dim]: + for causal_x, causal_y in product([True, False], [True, False]): + is_causal = (causal_x, causal_y) + self._test_all_dtypes_against_reference( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + head_dim_v=head_dim_v, + input_shape=input_shape, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + ) + + @skip_if_not_running_extended_tests() + @skip_if_libnatten_is_not_supported() + def test_2d_against_reference_extended(self): + problem_sizes = [ + (1, 2, 2, 32, (5, 25), (5, 16), (1, 16), (1, 1)), + (1, 2, 1, 32, (5, 25), (5, 16), (1, 16), (1, 1)), + (1, 2, 2, 32, (5, 25), (5, 16), (4, 16), (1, 1)), + (1, 2, 1, 32, (5, 25), (5, 16), (4, 16), (1, 1)), + (1, 1, 1, 32, (84, 69), (7, 68), (1, 1), (5, 1)), + (1, 1, 1, 32, (84, 69), (7, 23), (1, 1), (5, 1)), + (1, 1, 1, 32, (84, 69), (7, 20), (1, 6), (5, 1)), + (1, 1, 1, 128, (128, 128), (8, 8), (1, 1), (1, 1)), + (1, 1, 1, 128, (128, 128), (8, 8), (1, 1), (4, 4)), + (1, 1, 1, 128, (64, 64), (32, 32), (1, 1), (2, 2)), + (1, 1, 1, 128, (64, 64), (32, 32), (1, 1), (1, 2)), + (1, 1, 1, 128, (48, 48), (24, 24), (1, 1), (1, 2)), + (1, 1, 1, 128, (17, 48), (16, 24), (1, 1), (1, 2)), + (1, 1, 1, 128, (48, 48), (24, 24), (1, 1), (2, 2)), + (1, 1, 1, 128, (48, 17), (24, 16), (1, 1), (1, 1)), + (1, 1, 1, 128, (72, 80), (24, 16), (1, 1), (3, 5)), + (1, 1, 1, 128, (48, 17), (24, 16), (1, 1), (1, 1)), + (1, 1, 1, 128, (44, 80), (32, 32), (22, 16), (1, 1)), + (1, 1, 1, 128, (44, 80), (24, 24), (1, 1), (1, 1)), + (1, 1, 1, 128, (44, 80), (24, 24), (8, 16), (1, 1)), + (1, 1, 1, 128, (44, 80), (24, 16), (1, 1), (1, 1)), + (1, 1, 1, 128, (44, 80), (24, 16), (8, 16), (1, 1)), + (1, 1, 1, 64, (28, 40), (28, 40), (1, 1), (1, 1)), + (1, 1, 1, 32, (16, 16), (16, 16), (4, 5), (1, 1)), + (1, 1, 1, 32, (16, 16), (16, 16), (4, 8), (1, 1)), + (1, 1, 1, 32, (16, 16), (15, 15), (1, 1), (1, 1)), + (1, 1, 1, 32, (16, 16), (14, 14), (1, 1), (1, 1)), + (1, 1, 1, 64, (36, 40), (36, 40), (1, 1), (1, 1)), + (1, 1, 1, 128, (48, 80), (24, 24), (8, 8), (1, 1)), + ] + for i, ( + batch, + heads, + heads_kv, + head_dim, + input_shape, + kernel_size, + stride, + dilation, + ) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for causal_x, causal_y in product([True, False], [True, False]): + is_causal = (causal_x, causal_y) + self._test_all_dtypes_against_reference( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + head_dim_v=head_dim, + input_shape=input_shape, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + ) + + @skip_if_libnatten_is_not_supported() + def test_3d_against_reference(self): + problem_sizes = [ + (4, 8, 8, 64, (32, 10, 10), (7, 3, 3), (5, 1, 1), (1, 2, 3)), + (1, 4, 4, 32, (8, 8, 16), (3, 3, 3), (2, 1, 2), (2, 2, 4)), + (1, 4, 1, 32, (8, 8, 16), (3, 3, 3), (2, 1, 2), (2, 2, 4)), + (1, 1, 1, 64, (18, 37, 12), (14, 16, 12), (12, 8, 1), (1, 2, 1)), + (1, 1, 1, 32, (13, 11, 9), (3, 4, 3), (2, 3, 3), (3, 2, 2)), + (1, 1, 1, 32, (16, 16, 16), (16, 16, 16), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 32, (8, 4, 8), (7, 3, 7), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 32, (16, 16, 16), (16, 16, 16), (2, 4, 5), (1, 1, 1)), + ] + for i, ( + batch, + heads, + heads_kv, + head_dim, + input_shape, + kernel_size, + stride, + dilation, + ) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for head_dim_v in [random.choice(range(8, 193, 8)), head_dim]: + for causal_x, causal_y, causal_z in product( + [True, False], [True, False], [True, False] + ): + is_causal = (causal_x, causal_y, causal_z) + self._test_all_dtypes_against_reference( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + head_dim_v=head_dim_v, + input_shape=input_shape, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + ) + + @skip_if_not_running_extended_tests() + @skip_if_libnatten_is_not_supported() + def test_3d_against_reference_extended(self): + # GQA/MQA is not tested here -- + # Error rate is high for this backend due to manual (graph transform) implementation + # Combined with relatively large inputs frequent mismatches by a few elements is unavoidable + problem_sizes = [ + (1, 12, 64, (14, 8, 8), (7, 5, 5), (2, 1, 3), (2, 1, 1)), + (1, 1, 16, (18, 37, 12), (14, 16, 12), (12, 8, 1), (1, 2, 1)), + # (1, 1, 16, (32, 64, 64), (16, 16, 16), (1, 2, 2), (1, 1, 1)), # seqlen=131072 + # (1, 1, 16, (30, 48, 80), (18, 24, 24), (16, 8, 8), (1, 1, 1)), # seqlen=115200 + (1, 1, 128, (16, 44, 80), (12, 32, 32), (8, 22, 16), (1, 1, 1)), + (1, 1, 128, (31, 32, 32), (10, 32, 32), (1, 1, 1), (1, 1, 1)), + # (1, 1, 128, (57, 20, 88), (20, 4, 6), (1, 1, 1), (2, 1, 1)), # seqlen=100320 + (1, 1, 128, (57, 32, 32), (10, 32, 32), (1, 1, 1), (1, 1, 1)), + (1, 1, 128, (30, 32, 32), (10, 32, 32), (1, 1, 1), (1, 1, 1)), + (1, 1, 32, (16, 16, 16), (15, 15, 15), (1, 1, 1), (1, 1, 1)), + (1, 1, 64, (18, 37, 12), (14, 16, 12), (12, 8, 1), (1, 1, 1)), + (1, 1, 32, (13, 11, 9), (3, 4, 3), (1, 1, 1), (3, 2, 2)), + (1, 1, 32, (8, 8, 4), (3, 4, 3), (1, 1, 1), (1, 1, 1)), + (2, 2, 32, (8, 8, 10), (3, 4, 3), (3, 4, 1), (1, 1, 1)), + (1, 2, 32, (8, 8, 12), (5, 8, 11), (2, 3, 4), (1, 1, 1)), + (1, 1, 64, (18, 37, 12), (14, 16, 12), (12, 8, 1), (1, 1, 1)), + (2, 3, 32, (5, 45, 73), (2, 7, 32), (1, 4, 32), (2, 2, 2)), + # (1, 1, 128, (57, 20, 88), (19, 4, 6), (13, 1, 6), (2, 2, 7)), # seqlen=100320 + # (1, 1, 128, (57, 20, 88), (19, 4, 6), (13, 1, 6), (2, 2, 7)), # seqlen=100320, was batch=2 heads=3 + # (1, 1, 128, (32, 64, 64), (16, 16, 16), (2, 1, 1), (2, 2, 3)), # seqlen=131072 + # (1, 1, 128, (61, 61, 61), (10, 10, 10), (1, 1, 1), (1, 1, 2)), # seqlen=226981 + # (1, 1, 128, (61, 61, 61), (10, 10, 10), (1, 1, 1), (1, 2, 1)), # seqlen=226981 + # (1, 1, 128, (61, 61, 61), (10, 10, 10), (1, 1, 1), (2, 2, 2)), # seqlen=226981 + # (1, 1, 128, (32, 64, 64), (16, 16, 16), (2, 1, 1), (1, 1, 1)), # seqlen=131072 + # (1, 1, 128, (32, 64, 64), (16, 16, 16), (1, 1, 2), (1, 1, 1)), # seqlen=131072 + (1, 1, 32, (16, 16, 16), (16, 16, 16), (8, 4, 8), (1, 1, 1)), + (1, 1, 64, (24, 36, 40), (24, 36, 40), (10, 12, 13), (1, 1, 1)), + (1, 1, 128, (16, 44, 80), (16, 44, 80), (8, 4, 8), (1, 1, 1)), + # (1, 1, 128, (30, 44, 80), (18, 24, 24), (1, 1, 1), (1, 1, 1)), # seqlen=105600 + # (1, 1, 128, (30, 44, 80), (18, 24, 24), (2, 8, 16), (1, 1, 1)), # seqlen=105600 + # (1, 1, 128, (30, 44, 80), (24, 24, 16), (1, 1, 1), (1, 1, 1)), # seqlen=105600 + # (1, 1, 128, (30, 44, 80), (24, 24, 16), (2, 8, 16), (1, 1, 1)), # seqlen=105600 + (1, 1, 32, (16, 16, 16), (14, 14, 14), (1, 1, 1), (1, 1, 1)), + (1, 1, 32, (16, 16, 16), (3, 3, 3), (1, 1, 1), (1, 1, 1)), + (1, 1, 128, (16, 44, 80), (8, 9, 10), (1, 1, 1), (1, 1, 1)), + (1, 1, 64, (24, 28, 40), (11, 17, 31), (1, 1, 1), (1, 1, 1)), + (1, 1, 64, (24, 36, 40), (24, 36, 40), (1, 1, 1), (1, 1, 1)), + (1, 1, 128, (16, 44, 80), (16, 44, 80), (1, 1, 1), (1, 1, 1)), + (1, 1, 128, (30, 48, 18), (30, 48, 18), (1, 1, 1), (1, 1, 1)), + (1, 1, 128, (8, 8, 8), (8, 8, 8), (1, 1, 1), (1, 1, 1)), + (1, 1, 128, (30, 48, 17), (18, 24, 16), (1, 1, 1), (1, 1, 1)), + (1, 1, 128, (33, 33, 33), (18, 24, 16), (1, 1, 1), (1, 1, 1)), + (1, 1, 128, (30, 48, 18), (3, 2, 8), (1, 1, 1), (8, 18, 2)), + (1, 1, 128, (61, 32, 32), (10, 32, 32), (1, 1, 1), (2, 1, 1)), + (1, 1, 128, (57, 32, 32), (10, 32, 32), (1, 1, 1), (2, 1, 1)), + # (1, 1, 128, (57, 20, 88), (20, 8, 16), (1, 1, 1), (2, 1, 1)), # seqlen=100320 + (3, 1, 64, (18, 37, 12), (14, 16, 12), (12, 8, 6), (1, 2, 1)), + # (1, 1, 128, (32, 64, 64), (16, 16, 16), (1, 1, 2), (1, 3, 2)), # seqlen=131072 + # (1, 1, 128, (48, 64, 64), (7, 15, 11), (1, 2, 2), (5, 3, 2)), # seqlen=196608 + ] + for i, ( + batch, + heads, + head_dim, + input_shape, + kernel_size, + stride, + dilation, + ) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for causal_x, causal_y, causal_z in product( + [True, False], [True, False], [True, False] + ): + is_causal = (causal_x, causal_y, causal_z) + self._test_all_dtypes_against_reference( + batch=batch, + heads=heads, + head_dim=head_dim, + head_dim_v=head_dim, + input_shape=input_shape, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + ) + + def _test_randsweep_against_reference( + self, + na_dim, + max_tests=1000, + configs_to_test=10, + quick=False, + ): + max_seqlen = 2**17 if not quick else 2**16 + + # max size per-dim for different profiles + max_size = {1: 2**15, 2: 128, 3: 64} + + # seqlen limit for freely choosing batch and heads + seqlen_limit_batched = 2**13 + + for i in range(max_tests): + # to help with reproducibility of use cases + _reset_everything(random_seed=i, torch_seed=i) + + input_shape = [] + for j in range(na_dim): + input_shape.append(random.choice(range(4, max_size[na_dim] + 1))) + + while math.prod(input_shape) > max_seqlen: + dim_to_cut = random.choice(range(na_dim)) + input_shape[dim_to_cut] = max(4, int(input_shape[dim_to_cut] * 0.1)) + + input_shape = tuple(input_shape) + assert math.prod(input_shape) <= max_seqlen + + max_heads = min( + max(1, (seqlen_limit_batched // math.prod(input_shape)) * 4), 4 + ) + heads = random.choice(range(1, max_heads + 1)) + + max_batch = min( + max(1, ((seqlen_limit_batched // math.prod(input_shape)) * 4) // heads), + 4, + ) + batch = random.choice(range(1, max_batch + 1)) + + heads_kv = random.choice([i for i in range(1, heads + 1) if heads % i == 0]) + head_dim = random.choice(range(8, 193, 8)) + head_dim_v = random.choice( + range(max(8, head_dim - 16), min(193, head_dim + 16), 8) + ) + + kernel_size = tuple(random.choice(range(2, x + 1)) for x in input_shape) + stride = tuple(random.choice(range(1, k + 1)) for k in kernel_size) + dilation = tuple( + random.choice(range(1, x // k + 1)) + for x, k in zip(input_shape, kernel_size) + ) + is_causal = tuple(random.choice([False, True]) for _ in range(na_dim)) + + self._test_all_dtypes_against_reference( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + head_dim_v=head_dim_v, + input_shape=input_shape, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + configs_to_test=configs_to_test, + ) + + @skip_if_libnatten_is_not_supported() + def test_randsweep_1d_against_reference_quick(self): + self._test_randsweep_against_reference( + 1, max_tests=10, configs_to_test=5, quick=True + ) + + @skip_if_libnatten_is_not_supported() + def test_randsweep_2d_against_reference_quick(self): + self._test_randsweep_against_reference( + 2, max_tests=10, configs_to_test=3, quick=True + ) + + @skip_if_libnatten_is_not_supported() + def test_randsweep_3d_against_reference_quick(self): + self._test_randsweep_against_reference( + 3, max_tests=10, configs_to_test=3, quick=True + ) + + @skip_if_not_running_extended_tests() + @skip_if_libnatten_is_not_supported() + def test_randsweep_1d_against_reference(self): + self._test_randsweep_against_reference(1, max_tests=RAND_SWEEP_TESTS) + + @skip_if_not_running_extended_tests() + @skip_if_libnatten_is_not_supported() + def test_randsweep_2d_against_reference(self): + self._test_randsweep_against_reference( + 2, max_tests=RAND_SWEEP_TESTS, configs_to_test=3 + ) + + @skip_if_not_running_extended_tests() + @skip_if_libnatten_is_not_supported() + def test_randsweep_3d_against_reference(self): + self._test_randsweep_against_reference( + 3, max_tests=RAND_SWEEP_TESTS, configs_to_test=2 + ) + + def _test_determinism( + self, + batch, + heads, + head_dim, + input_shape, + kernel_size, + stride, + dilation, + is_causal=None, + configs_to_test=5, + fail_case=False, + ): + torch.set_default_device("cuda") + assert isinstance(input_shape, tuple) + + torch.use_deterministic_algorithms(not fail_case) + + ALLOWED_DTYPES = [ + torch.float32, + torch.float16, + torch.bfloat16, + ] + + for dtype in ALLOWED_DTYPES: + dummy = torch.randn( + (batch, *input_shape, heads, head_dim), + device="cuda", + dtype=dtype, + ) + + forward_configs = get_all_forward_configs(dummy) + backward_configs = get_all_backward_configs(dummy) + assert len(forward_configs) > 0 + assert len(backward_configs) > 0 + + random.shuffle(forward_configs) + random.shuffle(backward_configs) + + n_configs_to_test = min( + configs_to_test, max(len(forward_configs), len(backward_configs)) + ) + + for i in range(n_configs_to_test): + q_tile_shape, kv_tile_shape = forward_configs[i % len(forward_configs)] + backward_q_tile_shape, backward_kv_tile_shape = backward_configs[ + i % len(backward_configs) + ] + + tester = NattenBackendTester( + batch=batch, + heads=heads, + head_dim=head_dim, + head_dim_v=head_dim, + input_shape=input_shape, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + test_backprop=True, + reference_backend="cutlass-fna", + reference_fmha_backend="cutlass-fmha", + dtype=dtype, + reference_q_tile_shape=q_tile_shape, + reference_kv_tile_shape=kv_tile_shape, + reference_backward_q_tile_shape=backward_q_tile_shape, + reference_backward_kv_tile_shape=backward_kv_tile_shape, + ) + + tester.test( + eps=(0, (0, 0, 0)), + dtype=dtype, + target_backend="cutlass-fna", + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + backward_q_tile_shape=backward_q_tile_shape, + backward_kv_tile_shape=backward_kv_tile_shape, + ) + + torch.use_deterministic_algorithms(False) + + @skip_if_libnatten_is_not_supported() + def test_1d_determinism(self): + problem_sizes = [ + (2, 4, 128, (128,), (63,), (31,), (1,)), + (1, 1, 32, (128,), (3,), (2,), (5,)), + (1, 1, 64, (128,), (8,), (7,), (5,)), + ] + for i, ( + batch, + heads, + head_dim, + input_shape, + kernel_size, + stride, + dilation, + ) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for causal in [False, True]: + is_causal = (causal,) + self._test_determinism( + batch=batch, + heads=heads, + head_dim=head_dim, + input_shape=input_shape, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + ) + + @skip_if_libnatten_is_not_supported() + def test_2d_determinism(self): + problem_sizes = [ + (1, 1, 32, (84, 19), (7, 3), (1, 1), (5, 1)), + (1, 1, 128, (19, 29), (8, 8), (1, 1), (2, 3)), + (1, 1, 128, (48, 17), (24, 16), (1, 1), (2, 1)), + ] + for i, ( + batch, + heads, + head_dim, + input_shape, + kernel_size, + stride, + dilation, + ) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for causal_x, causal_y in product([False, True], [False, True]): + is_causal = (causal_x, causal_y) + self._test_determinism( + batch=batch, + heads=heads, + head_dim=head_dim, + input_shape=input_shape, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + ) + + @skip_if_libnatten_is_not_supported() + def test_3d_determinism(self): + problem_sizes = [ + (1, 2, 64, (14, 8, 8), (7, 5, 5), (2, 1, 3), (2, 1, 1)), + (3, 1, 16, (18, 37, 12), (14, 16, 12), (12, 8, 1), (1, 2, 1)), + (1, 1, 32, (16, 16, 16), (3, 3, 3), (1, 1, 1), (1, 1, 1)), + ] + for i, ( + batch, + heads, + head_dim, + input_shape, + kernel_size, + stride, + dilation, + ) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for causal_x, causal_y, causal_z in product( + [False, True], [False, True], [False, True] + ): + is_causal = (causal_x, causal_y, causal_z) + self._test_determinism( + batch=batch, + heads=heads, + head_dim=head_dim, + input_shape=input_shape, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + ) + + @pytest.mark.xfail( + reason="exact match must fail when determinism is disabled.", strict=True + ) + @skip_if_libnatten_is_not_supported() + def test_1d_determinism_xfail(self): + self._test_determinism( + batch=2, + heads=4, + head_dim=128, + input_shape=(128,), + kernel_size=(63,), + stride=(31,), + dilation=(1,), + is_causal=(False,), + fail_case=True, + ) + + @pytest.mark.xfail( + reason="exact match must fail when determinism is disabled.", strict=True + ) + @skip_if_libnatten_is_not_supported() + def test_2d_determinism_xfail(self): + self._test_determinism( + batch=1, + heads=1, + head_dim=128, + input_shape=(48, 17), + kernel_size=(24, 16), + stride=(1, 1), + dilation=(2, 1), + is_causal=(False, False), + fail_case=True, + ) + + @pytest.mark.xfail( + reason="exact match must fail when determinism is disabled.", strict=True + ) + @skip_if_libnatten_is_not_supported() + def test_3d_determinism_xfail(self): + self._test_determinism( + batch=1, + heads=12, + head_dim=64, + input_shape=(14, 8, 8), + kernel_size=(7, 5, 5), + stride=(2, 1, 3), + dilation=(2, 1, 1), + is_causal=(False, False, False), + fail_case=True, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/natten/tests/test_hopper_fna.py b/natten/tests/test_hopper_fna.py new file mode 100644 index 00000000..07277455 --- /dev/null +++ b/natten/tests/test_hopper_fna.py @@ -0,0 +1,549 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +import math +import random +import unittest +from itertools import product + +import torch +from natten._environment import _NUM_RAND_SWEEP_TESTS as RAND_SWEEP_TESTS +from natten.backends.configs.cutlass_hopper import ( + get_all_backward_configs, + get_all_forward_configs, +) +from natten.utils.testing import ( + skip_if_hopper_kernels_not_supported, + skip_if_libnatten_is_not_supported, + skip_if_not_running_extended_tests, +) + +from .utils import logger, NattenBackendTester + + +def _reset_everything(random_seed: int = 42, torch_seed: int = 42): + from natten.context import ( + NattenContext, + set_memory_usage_preference, + use_kv_parallelism_in_fused_na, + ) + + NattenContext.reset() + set_memory_usage_preference("unrestricted") + use_kv_parallelism_in_fused_na(True) + + random.seed(random_seed) + torch.manual_seed(torch_seed) + logger.debug(f"Reset seeds: {random_seed=}, {torch_seed=}") + torch.cuda.empty_cache() + torch.use_deterministic_algorithms(False) + + +class HopperFNABackendTest(unittest.TestCase): + def setUp(self): + _reset_everything() + + def tearDown(self): + _reset_everything() + + def _test_all_dtypes_against_cutlass_2x_fna( + self, + batch, + heads, + head_dim, + input_shape, + kernel_size, + stride, + dilation, + is_causal=None, + configs_to_test=None, + heads_kv=None, + ): + torch.set_default_device("cuda") + assert isinstance(input_shape, tuple) + na_dim = len(input_shape) + assert na_dim in [1, 2, 3], "Only supports NA1D, 2D, 3D." + + tester = NattenBackendTester( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + input_shape=input_shape, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + test_backprop=True, + reference_backend="cutlass-fna", + reference_fmha_backend="cutlass-fmha", + dtype=torch.float32, + ) + + ALLOWED_DTYPES = [ + (torch.float16, (1e-2, (3e-2, 3e-2, 3e-2))), + (torch.bfloat16, (5e-2, (3e-2, 3e-2, 3e-2))), + ] + + test_id = 0 + for dtype, atol in ALLOWED_DTYPES: + + dummy = torch.randn( + (batch, *input_shape, heads, head_dim), device="cuda", dtype=dtype + ) + forward_configs = get_all_forward_configs(dummy) + backward_configs = get_all_backward_configs(dummy) + assert len(forward_configs) > 0 + assert len(backward_configs) > 0 + + random.shuffle(forward_configs) + random.shuffle(backward_configs) + + for i in range(max(len(forward_configs), len(backward_configs))): + (q_tile_shape, kv_tile_shape), kernel_schedule = forward_configs[ + i % len(forward_configs) + ] + backward_q_tile_shape, backward_kv_tile_shape = backward_configs[ + i % len(backward_configs) + ] + + tester.test( + eps=atol, + dtype=dtype, + target_backend="hopper-fna", + target_fmha_backend="hopper-fmha", + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + backward_q_tile_shape=backward_q_tile_shape, + backward_kv_tile_shape=backward_kv_tile_shape, + kernel_schedule=kernel_schedule, + ) + test_id += 1 + if configs_to_test is not None and test_id > configs_to_test: + return + + @skip_if_libnatten_is_not_supported() + @skip_if_hopper_kernels_not_supported() + def test_1d_against_cutlass_2x(self): + problem_sizes = [ + (2, 4, 4, 128, (128,), (63,), (31,), (1)), + (2, 4, 2, 128, (128,), (63,), (31,), (1)), + (2, 4, 1, 128, (128,), (63,), (31,), (1)), + (4, 3, 3, 128, (256,), (255,), (82,), (1)), + (4, 3, 1, 128, (256,), (255,), (82,), (1)), + (1, 1, 1, 128, (128,), (8,), (1,), (1)), + (1, 1, 1, 128, (256,), (160,), (1,), (1)), + (1, 1, 1, 128, (256,), (150,), (1,), (1)), + (1, 1, 1, 128, (256,), (96,), (1,), (1)), + (1, 1, 1, 128, (256,), (8,), (1,), (1)), + (1, 1, 1, 128, (512,), (8,), (1,), (1)), + (1, 1, 1, 128, (32768,), (63,), (1,), (4,)), + (1, 1, 1, 128, (32768,), (63,), (1,), (8,)), + (1, 1, 1, 128, (32768,), (63,), (1,), (16,)), + (1, 1, 1, 128, (32768,), (2048,), (1,), (1)), + (1, 1, 1, 128, (32768,), (2048,), (128,), (1)), + (1, 1, 1, 128, (32768,), (2048,), (256,), (1)), + (1, 1, 1, 128, (32768,), (2048,), (2048,), (1)), + (1, 1, 1, 128, (64,), (8,), (1,), (1)), + (1, 2, 2, 128, (128,), (8,), (1,), (1)), + (1, 1, 1, 32, (64,), (8,), (1,), (1)), + (1, 2, 2, 32, (64,), (8,), (1,), (1)), + (1, 2, 2, 32, (64,), (3,), (1,), (1)), + (1, 2, 2, 32, (64,), (3,), (1,), (1)), + (1, 1, 1, 32, (64,), (3,), (1,), (1)), + (1, 1, 1, 32, (67,), (8,), (1,), (1)), + (1, 1, 1, 32, (69,), (12,), (1,), (1)), + (1, 1, 1, 32, (125,), (12,), (1,), (1)), + (1, 1, 1, 32, (17,), (12,), (1,), (1)), + (1, 1, 1, 32, (128,), (3,), (1,), (2)), + (1, 1, 1, 32, (128,), (3,), (2,), (5)), + (1, 1, 1, 128, (256,), (128,), (128,), (1)), + (1, 1, 1, 128, (256,), (136,), (136,), (1)), + (1, 1, 1, 128, (256,), (128,), (127,), (1)), + (1, 1, 1, 32, (256,), (4,), (1,), (1)), + (1, 1, 1, 128, (256,), (3,), (1,), (1)), + (1, 1, 1, 32, (256,), (64,), (1,), (1)), + (1, 1, 1, 32, (256,), (128,), (1,), (1)), + (1, 1, 1, 128, (256,), (128,), (128,), (1)), + (1, 1, 1, 32, (128,), (127,), (1,), (1)), + (1, 1, 1, 32, (128,), (3,), (1,), (1)), + (1, 2, 2, 64, (128,), (15,), (1,), (1)), + (1, 1, 1, 32, (128,), (3,), (2,), (10)), + (1, 1, 1, 64, (128,), (8,), (7,), (5)), + (1, 1, 1, 128, (128,), (61,), (33,), (1)), + (1, 1, 1, 32, (125,), (3,), (1,), (1)), + (1, 2, 2, 64, (125,), (15,), (1,), (1)), + (1, 1, 1, 128, (256,), (3,), (2,), (10)), + ] + for i, ( + batch, + heads, + heads_kv, + head_dim, + input_shape, + kernel_size, + stride, + dilation, + ) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for causal in [True, False]: + is_causal = (causal,) + self._test_all_dtypes_against_cutlass_2x_fna( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + input_shape=input_shape, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + ) + + @skip_if_libnatten_is_not_supported() + @skip_if_hopper_kernels_not_supported() + def test_2d_against_cutlass_2x(self): + problem_sizes = [ + (1, 2, 1, 32, (32, 16), (3, 3), (1, 1), (1, 1)), + (1, 3, 1, 32, (5, 25), (5, 16), (1, 16), (1, 1)), + (1, 1, 1, 32, (84, 19), (7, 3), (1, 1), (5, 1)), + (1, 1, 1, 128, (19, 29), (8, 8), (1, 1), (2, 3)), + (1, 1, 1, 128, (48, 17), (24, 16), (1, 1), (2, 1)), + (1, 1, 1, 128, (67, 80), (12, 7), (8, 4), (5, 11)), + (1, 1, 1, 128, (8, 8), (8, 8), (1, 1), (1, 1)), + (1, 1, 1, 128, (33, 33), (24, 16), (1, 1), (1, 1)), + (1, 1, 1, 32, (16, 16), (16, 16), (1, 1), (1, 1)), + (1, 1, 1, 128, (44, 80), (44, 80), (1, 1), (1, 1)), + (1, 1, 1, 32, (40, 20), (3, 7), (1, 1), (1, 1)), + (1, 1, 1, 32, (16, 16), (3, 3), (1, 1), (1, 1)), + (1, 1, 1, 128, (44, 80), (9, 10), (1, 1), (1, 1)), + (1, 1, 1, 64, (28, 40), (17, 31), (1, 1), (1, 1)), + (1, 1, 1, 64, (36, 40), (36, 40), (12, 13), (1, 1)), + (1, 1, 1, 128, (44, 80), (44, 80), (4, 8), (1, 1)), + ] + for i, ( + batch, + heads, + heads_kv, + head_dim, + input_shape, + kernel_size, + stride, + dilation, + ) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for causal_x, causal_y in product([False, True], [False, True]): + is_causal = (causal_x, causal_y) + self._test_all_dtypes_against_cutlass_2x_fna( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + input_shape=input_shape, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + ) + + @skip_if_not_running_extended_tests() + @skip_if_libnatten_is_not_supported() + @skip_if_hopper_kernels_not_supported() + def test_2d_against_cutlass_2x_extended(self): + problem_sizes = [ + (1, 2, 1, 32, (5, 25), (5, 16), (1, 16), (1, 1)), + (1, 2, 2, 32, (5, 25), (5, 16), (4, 16), (1, 1)), + (1, 1, 1, 32, (84, 69), (7, 68), (1, 1), (5, 1)), + (1, 1, 1, 32, (84, 69), (7, 23), (1, 1), (5, 1)), + (1, 1, 1, 32, (84, 69), (7, 20), (1, 6), (5, 1)), + (1, 1, 1, 128, (128, 128), (8, 8), (1, 1), (1, 1)), + (1, 1, 1, 128, (128, 128), (8, 8), (1, 1), (4, 4)), + (1, 1, 1, 128, (64, 64), (32, 32), (1, 1), (2, 2)), + (1, 1, 1, 128, (64, 64), (32, 32), (1, 1), (1, 2)), + (1, 1, 1, 128, (48, 48), (24, 24), (1, 1), (1, 2)), + (1, 1, 1, 128, (17, 48), (16, 24), (1, 1), (1, 2)), + (1, 1, 1, 128, (48, 48), (24, 24), (1, 1), (2, 2)), + (1, 1, 1, 128, (48, 17), (24, 16), (1, 1), (1, 1)), + (1, 1, 1, 128, (72, 80), (24, 16), (1, 1), (3, 5)), + (1, 1, 1, 128, (48, 17), (24, 16), (1, 1), (1, 1)), + (1, 1, 1, 128, (44, 80), (32, 32), (22, 16), (1, 1)), + (1, 1, 1, 128, (44, 80), (24, 24), (1, 1), (1, 1)), + (1, 1, 1, 128, (44, 80), (24, 24), (8, 16), (1, 1)), + (1, 1, 1, 128, (44, 80), (24, 16), (1, 1), (1, 1)), + (1, 1, 1, 128, (44, 80), (24, 16), (8, 16), (1, 1)), + (1, 1, 1, 64, (28, 40), (28, 40), (1, 1), (1, 1)), + (1, 1, 1, 32, (16, 16), (16, 16), (4, 5), (1, 1)), + (1, 1, 1, 32, (16, 16), (16, 16), (4, 8), (1, 1)), + (1, 1, 1, 32, (16, 16), (15, 15), (1, 1), (1, 1)), + (1, 1, 1, 32, (16, 16), (14, 14), (1, 1), (1, 1)), + (1, 1, 1, 64, (36, 40), (36, 40), (1, 1), (1, 1)), + (1, 1, 1, 128, (48, 80), (24, 24), (8, 8), (1, 1)), + ] + for i, ( + batch, + heads, + heads_kv, + head_dim, + input_shape, + kernel_size, + stride, + dilation, + ) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for causal_x, causal_y in product([True, False], [True, False]): + is_causal = (causal_x, causal_y) + self._test_all_dtypes_against_cutlass_2x_fna( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + input_shape=input_shape, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + ) + + @skip_if_libnatten_is_not_supported() + @skip_if_hopper_kernels_not_supported() + def test_3d_against_cutlass_2x(self): + problem_sizes = [ + (2, 2, 1, 128, (68, 64, 8), (14, 32, 2), (3, 4, 1), (1, 2, 3)), + (4, 8, 2, 64, (32, 10, 10), (7, 3, 3), (5, 1, 1), (1, 2, 3)), + (1, 1, 1, 64, (18, 37, 12), (14, 16, 12), (12, 8, 1), (1, 2, 1)), + (1, 1, 1, 32, (13, 11, 9), (3, 4, 3), (2, 3, 3), (3, 2, 2)), + (1, 4, 4, 32, (8, 8, 16), (3, 3, 3), (2, 1, 2), (2, 2, 4)), + (1, 1, 1, 64, (18, 37, 12), (14, 16, 12), (12, 8, 1), (1, 2, 1)), + (1, 1, 1, 128, (57, 20, 88), (20, 4, 6), (1, 1, 1), (2, 1, 1)), + (1, 1, 1, 128, (57, 32, 32), (10, 32, 32), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 128, (32, 64, 64), (16, 16, 16), (1, 2, 2), (1, 1, 1)), + (1, 1, 1, 128, (30, 48, 80), (18, 24, 24), (16, 8, 8), (1, 1, 1)), + (1, 1, 1, 128, (16, 44, 80), (12, 32, 32), (8, 22, 16), (1, 1, 1)), + (1, 1, 1, 128, (31, 32, 32), (10, 32, 32), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 128, (30, 32, 32), (10, 32, 32), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 32, (16, 16, 16), (16, 16, 16), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 32, (8, 4, 8), (7, 3, 7), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 32, (16, 16, 16), (15, 15, 15), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 64, (24, 28, 40), (24, 28, 40), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 32, (16, 16, 16), (16, 16, 16), (2, 4, 5), (1, 1, 1)), + ] + for i, ( + batch, + heads, + heads_kv, + head_dim, + input_shape, + kernel_size, + stride, + dilation, + ) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for causal_x, causal_y, causal_z in product( + [True, False], [True, False], [True, False] + ): + is_causal = (causal_x, causal_y, causal_z) + self._test_all_dtypes_against_cutlass_2x_fna( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + input_shape=input_shape, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + ) + + @skip_if_not_running_extended_tests() + @skip_if_libnatten_is_not_supported() + @skip_if_hopper_kernels_not_supported() + def test_3d_against_cutlass_2x_extended(self): + problem_sizes = [ + (1, 12, 3, 64, (14, 8, 8), (7, 5, 5), (2, 1, 3), (2, 1, 1)), + (1, 12, 4, 64, (14, 8, 8), (7, 5, 5), (2, 1, 3), (2, 1, 1)), + (1, 12, 6, 64, (14, 8, 8), (7, 5, 5), (2, 1, 3), (2, 1, 1)), + (1, 12, 12, 64, (32, 8, 8), (7, 5, 5), (2, 1, 3), (2, 1, 1)), + (1, 1, 1, 64, (18, 37, 12), (14, 16, 12), (12, 8, 1), (1, 1, 1)), + (1, 1, 1, 32, (13, 11, 9), (3, 4, 3), (1, 1, 1), (3, 2, 2)), + (1, 1, 1, 32, (8, 8, 4), (3, 4, 3), (1, 1, 1), (1, 1, 1)), + (2, 2, 2, 32, (8, 8, 10), (3, 4, 3), (3, 4, 1), (1, 1, 1)), + (1, 2, 2, 32, (8, 8, 12), (5, 8, 11), (2, 3, 4), (1, 1, 1)), + (1, 1, 1, 64, (18, 37, 12), (14, 16, 12), (12, 8, 1), (1, 1, 1)), + (2, 3, 3, 32, (5, 45, 73), (2, 7, 32), (1, 4, 32), (2, 2, 2)), + (1, 1, 1, 128, (57, 20, 88), (19, 4, 6), (13, 1, 6), (2, 2, 7)), + (2, 3, 3, 128, (57, 20, 88), (19, 4, 6), (13, 1, 6), (2, 2, 7)), + (1, 1, 1, 128, (32, 64, 64), (16, 16, 16), (2, 1, 1), (2, 2, 3)), + (1, 1, 1, 128, (61, 61, 61), (10, 10, 10), (1, 1, 1), (1, 1, 2)), + (1, 1, 1, 128, (61, 61, 61), (10, 10, 10), (1, 1, 1), (1, 2, 1)), + (1, 1, 1, 128, (61, 61, 61), (10, 10, 10), (1, 1, 1), (2, 2, 2)), + (1, 1, 1, 128, (32, 64, 64), (16, 16, 16), (2, 1, 1), (1, 1, 1)), + (1, 1, 1, 128, (32, 64, 64), (16, 16, 16), (1, 1, 2), (1, 1, 1)), + (1, 1, 1, 32, (16, 16, 16), (16, 16, 16), (8, 4, 8), (1, 1, 1)), + (1, 1, 1, 64, (24, 36, 40), (24, 36, 40), (10, 12, 13), (1, 1, 1)), + (1, 1, 1, 128, (16, 44, 80), (16, 44, 80), (8, 4, 8), (1, 1, 1)), + (1, 1, 1, 128, (30, 44, 80), (18, 24, 24), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 128, (30, 44, 80), (18, 24, 24), (2, 8, 16), (1, 1, 1)), + (1, 1, 1, 128, (30, 44, 80), (24, 24, 16), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 128, (30, 44, 80), (24, 24, 16), (2, 8, 16), (1, 1, 1)), + (1, 1, 1, 32, (16, 16, 16), (14, 14, 14), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 32, (16, 16, 16), (3, 3, 3), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 128, (16, 44, 80), (8, 9, 10), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 64, (24, 28, 40), (11, 17, 31), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 64, (24, 36, 40), (24, 36, 40), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 128, (16, 44, 80), (16, 44, 80), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 128, (30, 48, 18), (30, 48, 18), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 128, (8, 8, 8), (8, 8, 8), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 128, (30, 48, 17), (18, 24, 16), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 128, (33, 33, 33), (18, 24, 16), (1, 1, 1), (1, 1, 1)), + (1, 1, 1, 128, (30, 48, 18), (3, 2, 8), (1, 1, 1), (8, 18, 2)), + (1, 1, 1, 128, (61, 32, 32), (10, 32, 32), (1, 1, 1), (2, 1, 1)), + (1, 1, 1, 128, (57, 32, 32), (10, 32, 32), (1, 1, 1), (2, 1, 1)), + (1, 1, 1, 128, (57, 20, 88), (20, 8, 16), (1, 1, 1), (2, 1, 1)), + (3, 1, 1, 64, (18, 37, 12), (14, 16, 12), (12, 8, 6), (1, 2, 1)), + (1, 1, 1, 128, (32, 64, 64), (16, 16, 16), (1, 1, 2), (1, 3, 2)), + (1, 1, 1, 128, (48, 64, 64), (7, 15, 11), (1, 2, 2), (5, 3, 2)), + ] + for i, ( + batch, + heads, + heads_kv, + head_dim, + input_shape, + kernel_size, + stride, + dilation, + ) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for causal_x, causal_y, causal_z in product( + [True, False], [True, False], [True, False] + ): + is_causal = (causal_x, causal_y, causal_z) + self._test_all_dtypes_against_cutlass_2x_fna( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + input_shape=input_shape, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + ) + + def _test_randsweep_against_cutlass_2x( + self, na_dim, max_tests=1000, configs_to_test=None + ): + max_seqlen = 2**17 + # max size per-dim for different profiles + max_size = {1: 2**15, 2: 128, 3: 96} + + # seqlen limit for freely choosing batch and heads + seqlen_limit_batched = 2**13 + + for i in range(max_tests): + # to help with reproducibility of use cases + _reset_everything(random_seed=i, torch_seed=i) + + input_shape = [] + for j in range(na_dim): + input_shape.append(random.choice(range(4, max_size[na_dim] + 1))) + + while math.prod(input_shape) > max_seqlen: + dim_to_cut = random.choice(range(na_dim)) + input_shape[dim_to_cut] = max(4, int(input_shape[dim_to_cut] * 0.1)) + + input_shape = tuple(input_shape) + assert math.prod(input_shape) <= max_seqlen + + max_heads = min( + max(1, (seqlen_limit_batched // math.prod(input_shape)) * 4), 4 + ) + heads = random.choice(range(1, max_heads + 1)) + + max_batch = min( + max(1, ((seqlen_limit_batched // math.prod(input_shape)) * 4) // heads), + 4, + ) + batch = random.choice(range(1, max_batch + 1)) + + heads_kv = random.choice([i for i in range(1, heads + 1) if heads % i == 0]) + head_dim = random.choice([32, 64, 128]) + + input_shape = tuple(input_shape) + kernel_size = tuple(random.choice(range(2, x)) for x in input_shape) + stride = tuple(random.choice(range(1, k + 1)) for k in kernel_size) + dilation = tuple( + random.choice(range(1, x // k + 1)) + for x, k in zip(input_shape, kernel_size) + ) + is_causal = tuple(random.choice([False, True]) for _ in range(na_dim)) + + self._test_all_dtypes_against_cutlass_2x_fna( + batch=batch, + heads=heads, + heads_kv=heads_kv, + head_dim=head_dim, + input_shape=input_shape, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + configs_to_test=configs_to_test, + ) + + @skip_if_libnatten_is_not_supported() + @skip_if_hopper_kernels_not_supported() + def test_randsweep_1d_against_cutlass_2x_quick(self): + self._test_randsweep_against_cutlass_2x(1, max_tests=10, configs_to_test=3) + + @skip_if_libnatten_is_not_supported() + @skip_if_hopper_kernels_not_supported() + def test_randsweep_2d_against_cutlass_2x_quick(self): + self._test_randsweep_against_cutlass_2x(2, max_tests=10, configs_to_test=3) + + @skip_if_libnatten_is_not_supported() + @skip_if_hopper_kernels_not_supported() + def test_randsweep_3d_against_cutlass_2x_quick(self): + self._test_randsweep_against_cutlass_2x(3, max_tests=10, configs_to_test=3) + + @skip_if_not_running_extended_tests() + @skip_if_libnatten_is_not_supported() + @skip_if_hopper_kernels_not_supported() + def test_randsweep_1d_against_cutlass_2x(self): + self._test_randsweep_against_cutlass_2x(1, max_tests=RAND_SWEEP_TESTS) + + @skip_if_not_running_extended_tests() + @skip_if_libnatten_is_not_supported() + @skip_if_hopper_kernels_not_supported() + def test_randsweep_2d_against_cutlass_2x(self): + self._test_randsweep_against_cutlass_2x(2, max_tests=RAND_SWEEP_TESTS) + + @skip_if_not_running_extended_tests() + @skip_if_libnatten_is_not_supported() + @skip_if_hopper_kernels_not_supported() + def test_randsweep_3d_against_cutlass_2x(self): + self._test_randsweep_against_cutlass_2x(3, max_tests=RAND_SWEEP_TESTS) + + +if __name__ == "__main__": + unittest.main() diff --git a/natten/tests/test_kernels_ci.py b/natten/tests/test_kernels_ci.py new file mode 100644 index 00000000..78a6eb3a --- /dev/null +++ b/natten/tests/test_kernels_ci.py @@ -0,0 +1,110 @@ +# Fast parity checks for the kernels-community CI runner +# (`nix run .#ci-test`, which runs `pytest -m kernels_ci`). +# +# These compare the auto-selected fused backend for the GPU the CI runs on +# (cutlass-fna, hopper-fna, or blackwell-fna) against NATTEN's reference CUDA +# kernels, forward and backward. The vendored upstream suite (test_fna.py and +# friends) provides exhaustive coverage; this file is a smoke-level subset. + +import pytest +import torch + +import natten # noqa: F401 +from natten.backends.reference import reference_fna_generic +from natten.functional import na1d, na2d, na3d + +requires_cuda = pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA" +) + +NA_FUNCS = {1: na1d, 2: na2d, 3: na3d} + +CASES = [ + # (input_shape, kernel_size, stride, dilation, is_causal) + ((64,), (9,), (1,), (1,), (False,)), + ((64,), (13,), (2,), (2,), (True,)), + ((20, 20), (5, 5), (1, 1), (1, 1), (False, False)), + ((20, 20), (7, 5), (2, 1), (2, 2), (False, True)), + ((8, 10, 12), (3, 5, 5), (1, 1, 1), (1, 1, 1), (False, False, False)), + ((8, 10, 12), (5, 3, 3), (1, 2, 2), (1, 2, 2), (True, False, True)), +] + + +@requires_cuda +@pytest.mark.kernels_ci +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize( + "input_shape,kernel_size,stride,dilation,is_causal", CASES +) +def test_fused_matches_reference( + input_shape, kernel_size, stride, dilation, is_causal, dtype +): + torch.manual_seed(42) + na_dim = len(input_shape) + batch, heads, head_dim = 2, 2, 64 + device = "cuda" + + q = torch.randn( + (batch, *input_shape, heads, head_dim), + device=device, + dtype=dtype, + requires_grad=True, + ) + k = torch.randn_like(q, requires_grad=True) + v = torch.randn_like(q, requires_grad=True) + d_out = torch.randn_like(q) + + q_ref = q.detach().clone().requires_grad_(True) + k_ref = k.detach().clone().requires_grad_(True) + v_ref = v.detach().clone().requires_grad_(True) + + out = NA_FUNCS[na_dim]( + q, + k, + v, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + ) + out.backward(d_out) + + out_ref = reference_fna_generic( + q_ref, + k_ref, + v_ref, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + ) + out_ref.backward(d_out) + + atol, rtol = (1e-2, 1e-2) if dtype == torch.float16 else (2e-2, 2e-2) + torch.testing.assert_close(out, out_ref, atol=atol, rtol=rtol) + torch.testing.assert_close(q.grad, q_ref.grad, atol=atol, rtol=rtol) + torch.testing.assert_close(k.grad, k_ref.grad, atol=atol, rtol=rtol) + torch.testing.assert_close(v.grad, v_ref.grad, atol=atol, rtol=rtol) + + +@requires_cuda +@pytest.mark.kernels_ci +def test_ops_registered(): + from natten._ops import ops + + for name in ( + "na1d_forward", + "na2d_forward", + "na3d_forward", + "na1d_backward", + "na2d_backward", + "na3d_backward", + "reference_na2d_forward", + "hopper_na2d_forward", + "blackwell_na2d_forward", + "fmha_forward", + "token_permute_2d", + "token_unpermute_2d", + "compute_delta", + ): + assert hasattr(ops, name), f"op {name} is not registered" diff --git a/natten/tests/test_token_permute.py b/natten/tests/test_token_permute.py new file mode 100644 index 00000000..1d960855 --- /dev/null +++ b/natten/tests/test_token_permute.py @@ -0,0 +1,332 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +import random +import unittest + +import torch +from natten.token_permute import token_permute_operation, token_unpermute_operation +from natten.utils.testing import skip_if_cuda_is_not_supported, supports_float16 + +from .utils import logger + + +def _reset_everything(random_seed: int = 42, torch_seed: int = 42): + random.seed(random_seed) + torch.manual_seed(torch_seed) + logger.debug(f"Reset seeds: {random_seed=}, {torch_seed=}") + torch.use_deterministic_algorithms(False) + + +class TokenPermuteTest(unittest.TestCase): + def setUp(self): + _reset_everything() + + def tearDown(self): + _reset_everything() + + def _test_permute_unpermute_torch( + self, B, H, S, D, tile_shape, dilation, flip_dims, eps, dtype, device="cuda" + ): + with torch.no_grad(): + tensor = torch.randn((B, *S, H, D), device=device, dtype=dtype) + + logger.debug( + f"Testing token permute: {tensor.shape=}, {tile_shape=}, {dilation=}, {flip_dims=}, {dtype=}." + ) + + tensor_copy = tensor.clone() + tensor_permuted, qkv_shape, _ = token_permute_operation( + tensor_copy, + tile_shape=tile_shape, + dilation=dilation, + flip_tiled_dims=flip_dims, + use_torch=True, + ) + tensor_out = token_unpermute_operation( + tensor_permuted, + token_layout_shape=qkv_shape, + tile_shape=tile_shape, + dilation=dilation, + flip_tiled_dims=flip_dims, + use_torch=True, + ) + + torch.testing.assert_close(tensor, tensor_out, atol=eps, rtol=0) + + def test_permute_torch_cpu(self): + problem_sizes = [ + # 1D + (1, 1, (80,), 32, (8,), None), + (1, 1, (16,), 32, (2,), None), + (1, 1, (16,), 32, (8,), None), + (1, 1, (8,), 32, (8,), None), + (1, 1, (80,), 128, (20,), None), + (1, 1, (80,), 32, (8,), (7,)), + (1, 1, (16,), 32, (2,), (3,)), + (1, 1, (16,), 32, (8,), (3,)), + (1, 1, (8,), 32, (8,), (2,)), + (1, 1, (8,), 32, (8,), (4,)), + (1, 1, (80,), 128, (40,), (13,)), + # 2D + (1, 1, (48, 80), 32, (8, 8), None), + (1, 1, (16, 16), 32, (8, 2), None), + (1, 1, (16, 16), 32, (8, 2), None), + (1, 1, (44, 8), 32, (22, 8), None), + (1, 1, (48, 80), 128, (4, 16), None), + (1, 1, (48, 80), 32, (8, 8), (2, 8)), + (1, 1, (16, 16), 32, (8, 2), (2, 4)), + (1, 1, (16, 16), 32, (8, 2), (2, 1)), + (1, 1, (44, 8), 32, (22, 8), (4, 8)), + # 3D + (1, 1, (30, 48, 80), 32, (4, 8, 8), None), + (1, 1, (16, 16, 16), 32, (4, 8, 2), None), + (1, 1, (32, 16, 16), 32, (4, 8, 2), None), + (1, 1, (32, 44, 8), 32, (16, 22, 8), None), + (1, 1, (30, 48, 80), 128, (2, 4, 16), None), + (1, 1, (30, 48, 80), 32, (4, 8, 8), (2, 5, 4)), + (1, 1, (16, 16, 16), 32, (4, 8, 2), (1, 9, 3)), + ] + for i, (B, H, S, D, tile_shape, dilation) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for flip_dims in [True, False]: + self._test_permute_unpermute_torch( + B=B, + H=H, + S=S, + D=D, + tile_shape=tile_shape, + dilation=dilation, + flip_dims=flip_dims, + eps=1e-3, + dtype=torch.float32, + device="cpu", + ) + + @skip_if_cuda_is_not_supported() + def test_permute_torch_cuda(self): + torch.set_default_device("cuda") + problem_sizes = [ + # 1D + (1, 1, (128,), 1, (2,), None), + (1, 1, (80,), 32, (8,), None), + (1, 1, (16,), 32, (2,), None), + (1, 1, (16,), 32, (8,), None), + (1, 1, (8,), 32, (8,), None), + (1, 1, (80,), 128, (20,), None), + (1, 1, (80,), 128, (40,), None), + (1, 1, (80,), 32, (8,), (7,)), + (1, 1, (16,), 32, (2,), (3,)), + (1, 1, (16,), 32, (8,), (3,)), + (1, 1, (8,), 32, (8,), (2,)), + (1, 1, (8,), 32, (8,), (4,)), + (1, 1, (80,), 128, (20,), (13,)), + (1, 1, (80,), 128, (40,), (13,)), + # 2D + (1, 1, (48, 80), 32, (8, 8), None), + (1, 1, (16, 16), 32, (8, 2), None), + (1, 1, (16, 16), 32, (8, 2), None), + (1, 1, (44, 8), 32, (22, 8), None), + (1, 1, (48, 80), 128, (4, 16), None), + (1, 1, (48, 80), 128, (4, 16), None), + (1, 1, (48, 80), 32, (8, 8), (2, 8)), + (1, 1, (16, 16), 32, (8, 2), (2, 4)), + (1, 1, (16, 16), 32, (8, 2), (2, 1)), + (1, 1, (44, 8), 32, (22, 8), (4, 8)), + (1, 1, (48, 80), 128, (4, 16), (11, 7)), + (1, 1, (48, 80), 128, (4, 16), (5, 11)), + # 3D + (1, 1, (30, 48, 80), 32, (4, 8, 8), None), + (1, 1, (16, 16, 16), 32, (4, 8, 2), None), + (1, 1, (32, 16, 16), 32, (4, 8, 2), None), + (1, 1, (32, 44, 8), 32, (16, 22, 8), None), + (1, 1, (30, 48, 80), 128, (2, 4, 16), None), + (1, 1, (30, 48, 80), 32, (4, 8, 8), (2, 5, 4)), + (1, 1, (16, 16, 16), 32, (4, 8, 2), (1, 9, 3)), + (1, 1, (32, 16, 16), 32, (4, 8, 2), (7, 7, 7)), + (1, 1, (32, 44, 8), 32, (16, 22, 8), (13, 2, 6)), + (1, 1, (30, 48, 80), 128, (2, 4, 16), (7, 3, 1)), + (1, 1, (57, 32, 32), 128, (8, 4, 8), (2, 1, 1)), + (1, 1, (57, 32, 32), 128, (4, 8, 8), (2, 1, 1)), + (1, 1, (57, 32, 32), 128, (2, 8, 16), (2, 1, 1)), + (1, 1, (57, 32, 32), 128, (4, 4, 16), (2, 1, 1)), + (1, 1, (57, 20, 88), 128, (8, 4, 8), (2, 1, 1)), + (1, 1, (57, 20, 88), 128, (4, 8, 8), (2, 1, 1)), + (1, 1, (57, 20, 88), 128, (2, 8, 16), (2, 1, 1)), + (1, 1, (57, 20, 88), 128, (4, 4, 16), (2, 1, 1)), + (1, 1, (57, 20, 88), 128, (8, 4, 8), (2, 1, 1)), + (1, 1, (57, 20, 88), 128, (4, 8, 8), (2, 1, 1)), + (1, 1, (57, 20, 88), 128, (2, 8, 16), (2, 1, 1)), + (1, 1, (57, 20, 88), 128, (4, 4, 16), (2, 1, 1)), + ] + for i, (B, H, S, D, tile_shape, dilation) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for flip_dims in [True, False]: + self._test_permute_unpermute_torch( + B=B, + H=H, + S=S, + D=D, + tile_shape=tile_shape, + dilation=dilation, + flip_dims=flip_dims, + eps=1e-3, + dtype=( + torch.float16 + if supports_float16(torch.get_default_device()) + else torch.float32 + ), + device="cuda", + ) + + @skip_if_cuda_is_not_supported() + def test_permute_kernel_cuda(self): + torch.set_default_device("cuda") + problem_sizes = [ + # 1D + (1, 1, (7,), 1, (3,), None), + (1, 1, (7,), 4, (3,), None), + (1, 1, (19,), 4, (3,), None), + (1, 1, (19,), 4, (3,), (7,)), + (1, 1, (128,), 4, (2,), None), + (1, 1, (80,), 32, (8,), None), + (1, 1, (16,), 32, (2,), None), + (1, 1, (16,), 32, (8,), None), + (1, 1, (8,), 32, (8,), None), + (1, 1, (80,), 128, (20,), None), + (1, 1, (80,), 128, (40,), None), + (1, 1, (80,), 32, (8,), (7,)), + (1, 1, (16,), 32, (2,), (3,)), + (1, 1, (16,), 32, (8,), (3,)), + (1, 1, (8,), 32, (8,), (2,)), + (1, 1, (8,), 32, (8,), (4,)), + (1, 1, (80,), 128, (20,), (13,)), + (1, 1, (80,), 128, (40,), (13,)), + # 2D + (1, 1, (48, 80), 32, (8, 8), None), + (1, 1, (16, 16), 32, (8, 2), None), + (1, 1, (16, 16), 32, (8, 2), None), + (1, 1, (44, 8), 32, (22, 8), None), + (1, 1, (48, 80), 128, (4, 16), None), + (1, 1, (48, 80), 128, (4, 16), None), + (1, 1, (48, 80), 32, (8, 8), (2, 8)), + (1, 1, (16, 16), 32, (8, 2), (2, 4)), + (1, 1, (16, 16), 32, (8, 2), (2, 1)), + (1, 1, (44, 8), 32, (22, 8), (4, 8)), + (1, 1, (48, 80), 128, (4, 16), (11, 7)), + (1, 1, (48, 80), 128, (4, 16), (5, 11)), + # 3D + (1, 1, (30, 48, 80), 32, (4, 8, 8), None), + (1, 1, (16, 16, 16), 32, (4, 8, 2), None), + (1, 1, (32, 16, 16), 32, (4, 8, 2), None), + (1, 1, (32, 44, 8), 32, (16, 22, 8), None), + (1, 1, (30, 48, 80), 128, (2, 4, 16), None), + (1, 1, (30, 48, 80), 32, (4, 8, 8), (2, 5, 4)), + (1, 1, (16, 16, 16), 32, (4, 8, 2), (1, 9, 3)), + (1, 1, (32, 16, 16), 32, (4, 8, 2), (7, 7, 7)), + (1, 1, (32, 44, 8), 32, (16, 22, 8), (13, 2, 6)), + (1, 1, (30, 48, 80), 128, (2, 4, 16), (7, 3, 1)), + (1, 1, (57, 32, 32), 128, (8, 4, 8), (2, 1, 1)), + (1, 1, (57, 32, 32), 128, (4, 8, 8), (2, 1, 1)), + (1, 1, (57, 32, 32), 128, (2, 8, 16), (2, 1, 1)), + (1, 1, (57, 32, 32), 128, (4, 4, 16), (2, 1, 1)), + (1, 1, (57, 20, 88), 128, (8, 4, 8), (2, 1, 1)), + (1, 1, (57, 20, 88), 128, (4, 8, 8), (2, 1, 1)), + (1, 1, (57, 20, 88), 128, (2, 8, 16), (2, 1, 1)), + (1, 1, (57, 20, 88), 128, (4, 4, 16), (2, 1, 1)), + (1, 1, (57, 20, 88), 128, (8, 4, 8), (2, 1, 1)), + (1, 1, (57, 20, 88), 128, (4, 8, 8), (2, 1, 1)), + (1, 1, (57, 20, 88), 128, (2, 8, 16), (2, 1, 1)), + (1, 1, (57, 20, 88), 128, (4, 4, 16), (2, 1, 1)), + (1, 1, (57, 20, 88), 1, (4, 4, 16), (2, 1, 1)), + ] + for i, (B, H, S, D, tile_shape, dilation) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for flip_dims in [True, False]: + for dtype in [torch.float32, torch.float16, torch.bfloat16]: + self._test_permute_cuda_kernel( + B=B, + H=H, + S=S, + D=D, + tile_shape=tile_shape, + dilation=dilation, + flip_dims=flip_dims, + eps=1e-4, + dtype=dtype, + device="cuda", + ) + + def _test_permute_cuda_kernel( + self, B, H, S, D, tile_shape, dilation, flip_dims, eps, dtype, device="cuda" + ): + with torch.no_grad(): + tensor = torch.randn((B, *S, H, D), device=device, dtype=dtype) + + logger.debug( + f"Testing token permute kernel against torch reference: {tensor.shape=}, {tile_shape=}, {dilation=}, {flip_dims=}, {dtype=}." + ) + + tensor_copy = tensor.clone() + tensor_permuted_ref, qkv_shape_ref, _ = token_permute_operation( + tensor_copy, + tile_shape=tile_shape, + dilation=dilation, + flip_tiled_dims=flip_dims, + use_torch=True, + ) + tensor_out_ref = token_unpermute_operation( + tensor_permuted_ref, + token_layout_shape=qkv_shape_ref, + tile_shape=tile_shape, + dilation=dilation, + flip_tiled_dims=flip_dims, + use_torch=True, + ) + + tensor_copy_2 = tensor.clone() + tensor_permuted, qkv_shape, _ = token_permute_operation( + tensor_copy_2, + tile_shape=tile_shape, + dilation=dilation, + flip_tiled_dims=flip_dims, + use_torch=False, + ) + tensor_out = token_unpermute_operation( + tensor_permuted, + token_layout_shape=qkv_shape, + tile_shape=tile_shape, + dilation=dilation, + flip_tiled_dims=flip_dims, + use_torch=False, + ) + + torch.testing.assert_close( + tensor_permuted, tensor_permuted_ref, atol=eps, rtol=0 + ) + torch.testing.assert_close(tensor_out, tensor_out_ref, atol=eps, rtol=0) + torch.testing.assert_close(tensor_out, tensor, atol=eps, rtol=0) + + +if __name__ == "__main__": + unittest.main() diff --git a/natten/tests/test_torch_compile.py b/natten/tests/test_torch_compile.py new file mode 100644 index 00000000..afd2b091 --- /dev/null +++ b/natten/tests/test_torch_compile.py @@ -0,0 +1,575 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +import random +import unittest +from typing import Optional + +import torch +from natten.backends import get_compatible_backends, get_compatible_fmha_backends +from natten.functional import attention, neighborhood_attention_generic +from natten.utils.testing import skip_if_libnatten_is_not_supported +from natten.utils.varlen import generate_varlen_parameters +from torch import nn + +from .utils import logger, reset_torch_compile + + +def _reset_everything(random_seed: int = 42, torch_seed: int = 42): + reset_torch_compile(cache_size_limit=1, recompile_limit=1) + torch.cuda.empty_cache() + + random.seed(random_seed) + torch.manual_seed(torch_seed) + logger.debug(f"Reset seeds: {random_seed=}, {torch_seed=}") + torch.use_deterministic_algorithms(False) + + +class NABlock(nn.Module): + def __init__( + self, + embed_dim: int, + num_heads: int, + qkv_bias: bool = True, + qk_scale: Optional[float] = None, + proj_drop: float = 0.0, + ): + super().__init__() + self.embed_dim = embed_dim + self.num_heads = num_heads + self.head_dim = self.embed_dim // self.num_heads + self.scale = qk_scale or self.head_dim**-0.5 + + self.q = nn.Linear(self.embed_dim, self.embed_dim, bias=qkv_bias) + self.kv = nn.Linear(self.embed_dim, self.embed_dim * 2, bias=qkv_bias) + self.proj = nn.Linear(self.embed_dim, self.embed_dim) + + def forward( + self, + x: torch.Tensor, + additional_context: Optional[torch.Tensor] = None, + *args, + **kwargs, + ): + B, *input_shape, C = x.shape + na_dim = len(input_shape) + + permutation_q = [0] + [x + 1 for x in range(na_dim)] + [na_dim + 1, na_dim + 2] + permutation_kv = ( + [na_dim + 1, 0] + [x + 1 for x in range(na_dim)] + [na_dim + 2, na_dim + 3] + ) + + q = ( + self.q(x) + .reshape(B, *input_shape, self.num_heads, self.head_dim) + .permute(*permutation_q) + ) + kv = ( + self.kv(x) + .reshape(B, *input_shape, 2, self.num_heads, self.head_dim) + .permute(*permutation_kv) + ) + k, v = kv[0], kv[1] + + add_k, add_v = None, None + if additional_context is not None: + add_kv = ( + self.kv(additional_context) + .reshape(B, -1, 2, self.num_heads, self.head_dim) + .permute(2, 0, 1, 3, 4) + ) + add_k, add_v = add_kv[0], add_kv[1] + + x = neighborhood_attention_generic( # type: ignore + q, + k, + v, + *args, + additional_keys=add_k, + additional_values=add_v, + **kwargs, + ) + x = x.reshape(B, *input_shape, C) + + return self.proj(x) + + +class FMHABlock(nn.Module): + def __init__( + self, + embed_dim: int, + num_heads: int, + mlp_ratio: int, + qkv_bias: bool = True, + ): + super().__init__() + + self.embed_dim = embed_dim + self.mlp_ratio = mlp_ratio + self.mlp_dim = int(self.embed_dim * self.mlp_ratio) + self.num_heads = num_heads + self.head_dim = self.embed_dim // self.num_heads + self.scale = self.head_dim**-0.5 + + self.q = nn.Linear(self.embed_dim, self.embed_dim, bias=qkv_bias) + self.kv = nn.Linear(self.embed_dim, self.embed_dim * 2, bias=qkv_bias) + self.proj = nn.Linear(self.embed_dim, self.embed_dim) + + self.mlp = nn.Sequential( + nn.Linear(self.embed_dim, self.mlp_dim), + nn.GELU(), + nn.Linear(self.mlp_dim, embed_dim), + ) + + def forward(self, x: torch.Tensor, c: torch.Tensor, *args, **kwargs): + B, sQ, D = x.shape + B, sK, D = c.shape + q = self.q(x).reshape(B, sQ, self.num_heads, self.head_dim) + k, v = ( + self.kv(c) + .reshape(B, sK, 2, self.num_heads, self.head_dim) + .permute(2, 0, 1, 3, 4) + ) + + x0 = attention(q, k, v, *args, **kwargs) + assert isinstance(x0, torch.Tensor) + x0 = x0.reshape(B, sQ, D) + + return self.mlp(x0) + + +class Block(nn.Module): + def __init__( + self, + embed_dim: int, + num_heads: int, + mlp_ratio: int, + qkv_bias: bool = True, + ): + super().__init__() + + self.attn = NABlock( + embed_dim=embed_dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + ) + + self.mlp = nn.Sequential( + nn.Linear(embed_dim, int(embed_dim * mlp_ratio)), + nn.GELU(), + nn.Linear(int(embed_dim * mlp_ratio), embed_dim), + ) + + def forward( + self, + x: torch.Tensor, + additional_context: Optional[torch.Tensor] = None, + *args, + **kwargs, + ): + x0 = self.attn(x, additional_context=additional_context, *args, **kwargs) + return self.mlp(x0) + + +class TorchCompileTests(unittest.TestCase): + def _test_na_module( + self, + batch, + token_layout_shape, + num_heads, + head_dim, + kernel_size, + dilation, + stride, + is_causal, + atol, + additional_tokens: int = 0, + ): + na_dim = len(token_layout_shape) + + embed_dim = num_heads * head_dim + + device = "cuda" + dtype = torch.float16 + + dummy = torch.randn( + (batch, *token_layout_shape, num_heads, head_dim), + dtype=dtype, + device=device, + requires_grad=True, + ) + # TODO: Flex FNA is also weirdly incompatible with torch compiling the model -- + # likely something to do with our mask... + backends = [ + b + for b in get_compatible_backends(dummy, dummy, dummy, torch_compile=False) + if b != "flex-fna" + ] + for backend in backends: + reset_torch_compile(cache_size_limit=1, recompile_limit=1) + + logger.debug( + f"Testing torch compile on {na_dim}-D module with input shapes: " + f"{batch=}, {num_heads=}, {head_dim=}, {token_layout_shape=}, " + f"{kernel_size=}, {dilation=}, {stride=}, {is_causal=}, " + f"{additional_tokens=}, {dtype=}, {device=}, {backend=}." + ) + + model = ( + Block( + embed_dim=embed_dim, + mlp_ratio=2, + num_heads=num_heads, + ) + .to(dtype) + .to(device) + ) + + forward_args = { + "kernel_size": kernel_size, + "dilation": dilation, + "stride": stride, + "is_causal": is_causal, + "backend": backend, + "torch_compile": True, + } + + model_eager = model + + model_compiled = torch.compile(model, fullgraph=True, backend="inductor") + + x = torch.randn( + (batch, *token_layout_shape, embed_dim), dtype=dtype, device=device + ) + dy = ( + torch.randn( + (batch, *token_layout_shape, embed_dim), dtype=dtype, device=device + ) + * 0.1 + ) + + x_ref = x.clone().requires_grad_(True) + dy_ref = dy.clone() + + additional_context, additional_context_ref = None, None + if additional_tokens > 0: + additional_context = torch.randn( + (batch, additional_tokens, embed_dim), dtype=dtype, device=device + ) + additional_context_ref = additional_context.clone().requires_grad_(True) + additional_context = additional_context.requires_grad_(True) + + # eager + y_ref = model_eager(x_ref, additional_context_ref, **forward_args) + y_ref.backward(dy_ref) + dx_ref = x_ref.grad + + # compile on first attempt + x = x.requires_grad_(True) + y = model_compiled(x, additional_context, **forward_args) + y.backward(dy) + dx = x.grad + + torch.testing.assert_close(y, y_ref, atol=atol, rtol=0) + torch.testing.assert_close(dx, dx_ref, atol=atol, rtol=0) + if additional_tokens > 0: + assert additional_context is not None + assert additional_context_ref is not None + assert additional_context.grad is not None + assert additional_context_ref.grad is not None + + dc = additional_context.grad + dc_ref = additional_context_ref.grad + torch.testing.assert_close(dc, dc_ref, atol=atol, rtol=0) + + # Second run, just to make sure it doesn't recompile or crash + y = model_compiled(x, additional_context, **forward_args) + y.backward(dy) + dx = x.grad + + def _test_fmha_module( + self, batch, seqlens_Q, seqlens_KV, num_heads, head_dim, is_causal, atol + ): + device = "cuda" + dtype = torch.float16 + + embed_dim = num_heads * head_dim + + assert len(seqlens_Q) == len(seqlens_KV) + assert len(seqlens_Q) >= 1 + assert len(seqlens_Q) == 1 or batch == len(seqlens_Q) + + seqlen_q = sum(seqlens_Q) + seqlen_kv = sum(seqlens_KV) + is_varlen = len(seqlens_Q) > 1 + + batch_ = 1 if is_varlen else batch + seqlens_Q_ = ( + torch.tensor(seqlens_Q, device=device, dtype=torch.int32) + if is_varlen + else None + ) + seqlens_KV_ = ( + torch.tensor(seqlens_KV, device=device, dtype=torch.int32) + if is_varlen + else None + ) + + dummy_q = torch.randn( + (batch_, seqlen_q, num_heads, head_dim), + dtype=dtype, + device=device, + requires_grad=True, + ) + dummy_kv = torch.randn( + (batch_, seqlen_kv, num_heads, head_dim), + dtype=dtype, + device=device, + requires_grad=True, + ) + + # TODO: Flex FMHA also fails... + backends = [ + b + for b in get_compatible_fmha_backends( + dummy_q, + dummy_kv, + dummy_kv, + is_causal=is_causal, + is_varlen=is_varlen, + torch_compile=True, + ) + if "b" != "flex-fmha" + ] + + # seq maxes MUST be computed ahead of time when using torch compile + # because need to be copied to host + ( + cumulative_seqlen_Q, + cumulative_seqlen_KV, + max_seqlen_Q, + max_seqlen_KV, + ) = generate_varlen_parameters( + query=dummy_q, + key=dummy_kv, + value=dummy_kv, + seqlens_Q=seqlens_Q_, + seqlens_KV=seqlens_KV_, + ) + + for backend in backends: + reset_torch_compile(cache_size_limit=1, recompile_limit=1) + + logger.debug( + f"Testing torch compile on FMHA module with input shapes: " + f"{batch=}, {num_heads=}, {head_dim=}, {seqlens_Q=}, {seqlens_KV=}, " + f"{is_causal=}, {dtype=}, {device=}, {backend=}." + ) + + model_eager = ( + FMHABlock( + embed_dim=embed_dim, + mlp_ratio=2, + num_heads=num_heads, + ) + .to(dtype) + .to(device) + ) + + model_compiled = torch.compile( + model_eager, fullgraph=True, backend="inductor" + ) + + x = torch.randn((batch_, seqlen_q, embed_dim), dtype=dtype, device=device) + c = torch.randn((batch_, seqlen_kv, embed_dim), dtype=dtype, device=device) + dy = ( + torch.randn((batch_, seqlen_q, embed_dim), dtype=dtype, device=device) + * 0.1 + ) + + x_ref = x.clone().requires_grad_(True) + c_ref = c.clone().requires_grad_(True) + dy_ref = dy.clone() + + # eager + y_ref = model_eager( + x_ref, + c_ref, + is_causal=is_causal, + cumulative_seqlen_Q=cumulative_seqlen_Q, + cumulative_seqlen_KV=cumulative_seqlen_KV, + max_seqlen_Q=max_seqlen_Q, + max_seqlen_KV=max_seqlen_KV, + backend=backend, + torch_compile=True, + ) + y_ref.backward(dy_ref) + dx_ref = x_ref.grad + dc_ref = c_ref.grad + + # compile on first attempt + x = x.requires_grad_(True) + c = c.requires_grad_(True) + y = model_compiled( + x, + c, + is_causal=is_causal, + cumulative_seqlen_Q=cumulative_seqlen_Q, + cumulative_seqlen_KV=cumulative_seqlen_KV, + max_seqlen_Q=max_seqlen_Q, + max_seqlen_KV=max_seqlen_KV, + backend=backend, + torch_compile=True, + ) + y.backward(dy) + dx = x.grad + dc = c.grad + + torch.testing.assert_close(y, y_ref, atol=atol, rtol=0) + torch.testing.assert_close(dx, dx_ref, atol=atol, rtol=0) + torch.testing.assert_close(dc, dc_ref, atol=atol, rtol=0) + + # Second run, just to make sure it doesn't recompile or crash + y = model_compiled( + x, + c, + is_causal=is_causal, + cumulative_seqlen_Q=cumulative_seqlen_Q, + cumulative_seqlen_KV=cumulative_seqlen_KV, + max_seqlen_Q=max_seqlen_Q, + max_seqlen_KV=max_seqlen_KV, + backend=backend, + torch_compile=True, + ) + y.backward(dy) + dx = x.grad + dc = c.grad + + @skip_if_libnatten_is_not_supported() + def test_compiled_na(self): + problems = [ + (2, (672,), 4, 32, (5,), (1,), (1,), (False,)), + (2, (688,), 4, 32, (5,), (1,), (1,), (False,)), + (2, (9215,), 4, 32, (111,), (1,), (1,), (False,)), + (2, (9215,), 4, 32, (9215,), (1,), (1,), (True,)), # fmha case + (2, (111,), 4, 32, (111,), (1,), (1,), (True,)), # fmha case + (2, (16384,), 4, 32, (128,), (16,), (8,), (False,)), + (2, (16384,), 4, 32, (128,), (16,), (8,), (True,)), + (1, (1024,), 2, 64, (64,), (2,), (2,), (False,)), + (1, (1024,), 2, 64, (64,), (2,), (2,), (True,)), + (8, (14, 14), 4, 32, (4, 4), (2, 2), (2, 3), (False, False)), + (8, (14, 14), 4, 32, (4, 4), (2, 2), (2, 3), (False, False)), + (2, (16, 24), 4, 32, (4, 6), (1, 1), (2, 1), (False, False)), + (2, (9215,), 4, 32, (111,), (2,), (3,), (False,)), + (2, (19, 27), 4, 32, (4, 2), (2, 7), (2, 1), (False, False)), + ( + 2, + (16, 24, 24), + 4, + 32, + (4, 6, 8), + (2, 3, 2), + (2, 1, 3), + (True, False, False), + ), + (8, (14, 14), 4, 32, (14, 14), (1, 1), (1, 1), (False, False)), # fmha case + (1, (16384,), 2, 128, (16384,), (1,), (1,), (False,)), # fmha case + (1, (16384,), 2, 128, (16384,), (1,), (1,), (True,)), # fmha case + (1, (3287,), 4, 64, (3287,), (1,), (1,), (False,)), # fmha case + (1, (3287,), 4, 64, (3287,), (1,), (1,), (True,)), # fmha case + ] + for i, problem in enumerate(problems): + _reset_everything(random_seed=i, torch_seed=i) + ( + batch, + token_layout_shape, + num_heads, + head_dim, + kernel_size, + dilation, + stride, + is_causal, + ) = problem + + for additional_tokens in [0, 64]: + self._test_na_module( + batch=batch, + token_layout_shape=token_layout_shape, + num_heads=num_heads, + head_dim=head_dim, + kernel_size=kernel_size, + dilation=dilation, + stride=stride, + is_causal=is_causal, + atol=1e-3, + additional_tokens=additional_tokens, + ) + + @skip_if_libnatten_is_not_supported() + def test_compiled_fmha(self): + problem_sizes = [ + (1, 4, 128, [128], [128]), + (1, 1, 128, [128], [1024]), + (1, 1, 128, [128], [13568]), + (1, 1, 128, [128], [13496]), + (1, 1, 32, [128], [13496]), + (1, 1, 32, [32], [13496]), + (3, 1, 32, [77], [8504]), + (1, 1, 32, [77], [8504]), + (1, 1, 64, [40], [12296]), + (1, 2, 64, [40], [12296]), + (1, 2, 64, [40], [12296]), + (1, 1, 128, [128], [128]), + (6, 1, 128, [128, 128, 135, 121, 128, 128], [128, 128, 135, 121, 128, 128]), + (5, 1, 128, [128, 128, 135, 128, 128], [128, 128, 135, 128, 128]), + (2, 1, 128, [135, 200], [128, 768]), + (2, 1, 128, [1024, 200], [128, 768]), + (2, 1, 128, [135, 200], [135, 768]), + (2, 1, 128, [1024, 200], [135, 768]), + (2, 1, 128, [1024, 256], [128, 768]), + (4, 1, 128, [1024, 8, 17, 2048], [10, 20, 512, 16]), + (3, 2, 128, [268, 1584, 1571], [2448, 4088, 1925]), + (2, 1, 128, [1024, 256], [512, 768]), + ] + for i, ( + batch, + num_heads, + head_dim, + seqlens_Q, + seqlens_KV, + ) in enumerate(problem_sizes): + _reset_everything(random_seed=i, torch_seed=i) + for is_causal in [False, True]: + self._test_fmha_module( + batch=batch, + seqlens_Q=seqlens_Q, + seqlens_KV=seqlens_KV, + num_heads=num_heads, + head_dim=head_dim, + is_causal=is_causal, + atol=1e-3, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/natten/tests/utils.py b/natten/tests/utils.py new file mode 100644 index 00000000..a61ac00e --- /dev/null +++ b/natten/tests/utils.py @@ -0,0 +1,319 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +import time +from typing import Optional, Tuple, Union + +import torch +from natten.backends.reference import reference_fna_generic +from natten.functional import neighborhood_attention_generic +from natten.types import CausalArgType, DimensionType, KernelSchedule +from natten.utils import log +from natten.utils.checks import check_all_args + +logger = log.get_logger("natten_tests") + + +def reset_torch_compile(cache_size_limit, recompile_limit: int | None = None): + recompile_limit = recompile_limit or cache_size_limit * 4 + # Torch compile reset and sensible settings for unit testing + logger.debug( + f"Resetting torch compile cache: {cache_size_limit=}, {recompile_limit=}." + ) + torch.compiler.reset() + torch._dynamo.config.cache_size_limit = cache_size_limit + torch._dynamo.config.accumulated_recompile_limit = recompile_limit + torch._dynamo.config.fail_on_recompile_limit_hit = True + + +# Runs one backend once as a reference, and may run another backend multiple times +# with different configurations. +class NattenBackendTester: + def __init__( + self, + batch: int, + heads: int, + head_dim: int, + input_shape: DimensionType, + kernel_size: DimensionType, + stride: DimensionType, + dilation: DimensionType, + is_causal: CausalArgType, + test_backprop: bool, + reference_backend: str, + reference_fmha_backend: str, + dtype: torch.dtype, + head_dim_v: Optional[int] = None, + heads_kv: Optional[int] = None, + reference_q_tile_shape: Optional[DimensionType] = None, + reference_kv_tile_shape: Optional[DimensionType] = None, + reference_backward_q_tile_shape: Optional[DimensionType] = None, + reference_backward_kv_tile_shape: Optional[DimensionType] = None, + reference_backward_kv_splits: Optional[DimensionType] = None, + reference_backward_use_pt_reduction: bool = False, + ): + assert isinstance(input_shape, tuple) + na_dim = len(input_shape) + assert na_dim in [1, 2, 3], "Only supports NA1D, 2D, 3D." + + self.batch = batch + self.heads = heads + self.heads_kv = heads_kv or heads + self.head_dim = head_dim + self.head_dim_v = head_dim_v or head_dim + self.input_shape = input_shape + self.kernel_size, self.stride, self.dilation, self.is_causal = check_all_args( + na_dim, kernel_size, stride, dilation, is_causal + ) + + self.test_backprop = test_backprop + self.reference_backend = reference_backend + self.reference_fmha_backend = reference_fmha_backend + + with torch.no_grad(): + orig_dtype = dtype + if dtype in [torch.float8_e5m2, torch.float8_e4m3fn]: + dtype = torch.float16 + + q_ref, k_ref, v_ref, d_out_ref = ( + torch.randn( + (self.batch, *self.input_shape, self.heads, self.head_dim), + device="cuda", + dtype=dtype, + ), + torch.randn( + (self.batch, *self.input_shape, self.heads_kv, self.head_dim), + device="cuda", + dtype=dtype, + ), + torch.randn( + (self.batch, *self.input_shape, self.heads_kv, self.head_dim_v), + device="cuda", + dtype=dtype, + ), + torch.randn( + (self.batch, *self.input_shape, self.heads, self.head_dim_v), + device="cuda", + dtype=dtype, + ) + * 0.05, + ) + + if dtype != orig_dtype: + q_ref = q_ref.to(orig_dtype) + k_ref = k_ref.to(orig_dtype) + v_ref = v_ref.to(orig_dtype) + d_out_ref = d_out_ref.to(orig_dtype) + + self.q, self.k, self.v, self.d_out = ( + q_ref.clone(), + k_ref.clone(), + v_ref.clone(), + d_out_ref.clone(), + ) + + # Reference + torch.cuda.synchronize() + start_time = time.time() + + q_ref.requires_grad_(True) + k_ref.requires_grad_(True) + v_ref.requires_grad_(True) + d_out_ref.requires_grad_(True) + if reference_backend is None or reference_backend == "reference": + out_ref_ = reference_fna_generic( + q_ref, + k_ref, + v_ref, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + return_lse=False, + ) + + else: + # TODO: don't rely on `neighborhood_attention_generic` finding the right backend + # and explicitly call the backend fns. + out_ref_ = neighborhood_attention_generic( + q_ref, + k_ref, + v_ref, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + backend=reference_backend, + q_tile_shape=reference_q_tile_shape, + kv_tile_shape=reference_kv_tile_shape, + backward_q_tile_shape=reference_backward_q_tile_shape, + backward_kv_tile_shape=reference_backward_kv_tile_shape, + backward_kv_splits=reference_backward_kv_splits, + backward_use_pt_reduction=reference_backward_use_pt_reduction, + attention_kwargs={ + "backend": reference_fmha_backend, + }, + ) + + self.out_ref = out_ref_.data.clone().float() # type: ignore[union-attr] + + self.dq_ref, self.dk_ref, self.dv_ref = None, None, None + if test_backprop: + out_ref_.backward(d_out_ref) # type: ignore[union-attr] + with torch.no_grad(): + assert q_ref.grad is not None + assert k_ref.grad is not None + assert v_ref.grad is not None + self.dq_ref, self.dk_ref, self.dv_ref = ( + q_ref.grad.clone().float(), + k_ref.grad.clone().float(), + v_ref.grad.clone().float(), + ) + + torch.cuda.synchronize() + reference_time = time.time() - start_time + logger.debug( + f"Reference ({reference_backend}/{reference_fmha_backend}) ran in {reference_time:.2f} seconds." + ) + + def test( + self, + eps: Union[ + float, Tuple[float, float], Tuple[float, Tuple[float, float, float]] + ], + dtype: torch.dtype, + target_backend: str, + target_fmha_backend: str = "cutlass-fmha", + q_tile_shape: Optional[DimensionType] = None, + kv_tile_shape: Optional[DimensionType] = None, + backward_q_tile_shape: Optional[DimensionType] = None, + backward_kv_tile_shape: Optional[DimensionType] = None, + backward_kv_splits: Optional[DimensionType] = None, + backward_use_pt_reduction: bool = False, + run_persistent_kernel: bool = True, + kernel_schedule: Optional[KernelSchedule] = None, + torch_compile: bool = False, + test_backprop: Optional[bool] = None, + ): + batch = self.batch + heads = self.heads + heads_kv = self.heads_kv + head_dim = self.head_dim + head_dim_v = self.head_dim_v + input_shape = self.input_shape + kernel_size = self.kernel_size + stride = self.stride + dilation = self.dilation + is_causal = self.is_causal + reference_backend = self.reference_backend + test_backprop_safe: bool = ( + self.test_backprop if test_backprop is None else test_backprop + ) + + logger.debug( + f"Testing {target_backend} against {reference_backend}:\n" + f"{batch=}, {heads=}, {heads_kv=}, {head_dim=}, {head_dim_v=}, {input_shape=}, {dtype=}\n" + f"{kernel_size=}, {stride=}, {dilation=}, {is_causal=},\n" + f"{q_tile_shape=}, {kv_tile_shape=}, {run_persistent_kernel=}, {kernel_schedule=}, " + f"{torch_compile=}" + + ( + f"\n{backward_q_tile_shape=}, {backward_kv_tile_shape=}, " + f"{backward_kv_splits=}, {backward_use_pt_reduction=}." + if test_backprop_safe + else "." + ) + ) + + q, k, v, d_out = ( + self.q.clone().to(dtype), + self.k.clone().to(dtype), + self.v.clone().to(dtype), + self.d_out.clone().to(dtype), + ) + q.requires_grad_(test_backprop_safe) + k.requires_grad_(test_backprop_safe) + v.requires_grad_(test_backprop_safe) + d_out.requires_grad_(test_backprop_safe) + + torch.cuda.synchronize() + start_time = time.time() + + out_: torch.Tensor = neighborhood_attention_generic( # type: ignore[assignment] + q, + k, + v, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + backend=target_backend, + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + backward_q_tile_shape=backward_q_tile_shape, + backward_kv_tile_shape=backward_kv_tile_shape, + backward_kv_splits=backward_kv_splits, + backward_use_pt_reduction=backward_use_pt_reduction, + run_persistent_kernel=run_persistent_kernel, + kernel_schedule=kernel_schedule, + torch_compile=torch_compile, + attention_kwargs={"backend": target_fmha_backend}, + ) + out = out_.data.clone().float() + + if test_backprop_safe: + dq, dk, dv = None, None, None + out_.backward(d_out) + with torch.no_grad(): + assert q.grad is not None + assert k.grad is not None + assert v.grad is not None + dq, dk, dv = ( + q.grad.clone().float(), + k.grad.clone().float(), + v.grad.clone().float(), + ) + + if isinstance(eps, tuple): + eps_forward, eps_backward = eps + else: + eps_forward, eps_backward = eps, eps + + torch.cuda.synchronize() + runtime = time.time() - start_time + logger.debug( + f"Backend ({target_backend}/{target_fmha_backend}) ran in {runtime:.2f} seconds." + ) + + torch.testing.assert_close(out, self.out_ref, atol=eps_forward, rtol=0) + + if test_backprop_safe: + if isinstance(eps_backward, tuple): + assert len(eps_backward) == 3 + eps_dq, eps_dk, eps_dv = eps_backward + else: + assert isinstance(eps_backward, float) + eps_dq, eps_dk, eps_dv = eps_backward, eps_backward, eps_backward + + torch.testing.assert_close(dq, self.dq_ref, atol=eps_dq, rtol=0) + torch.testing.assert_close(dk, self.dk_ref, atol=eps_dk, rtol=0) + torch.testing.assert_close(dv, self.dv_ref, atol=eps_dv, rtol=0) diff --git a/natten/torch-ext/natten/__init__.py b/natten/torch-ext/natten/__init__.py new file mode 100644 index 00000000..37ddeac9 --- /dev/null +++ b/natten/torch-ext/natten/__init__.py @@ -0,0 +1,106 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +from ._environment import HAS_LIBNATTEN +from .backends import ( + get_bwd_configs_for_cutlass_blackwell_fmha, + get_bwd_configs_for_cutlass_blackwell_fna, + get_bwd_configs_for_cutlass_fmha, + get_bwd_configs_for_cutlass_fna, + get_bwd_configs_for_cutlass_hopper_fmha, + get_bwd_configs_for_cutlass_hopper_fna, + get_configs_for_cutlass_blackwell_fmha, + get_configs_for_cutlass_blackwell_fna, + get_configs_for_cutlass_fmha, + get_configs_for_cutlass_fna, + get_configs_for_cutlass_hopper_fmha, + get_configs_for_cutlass_hopper_fna, + get_configs_for_flex_fmha, + get_configs_for_flex_fna, +) +from .context import ( + allow_flex_compile, + allow_flex_compile_backprop, + are_deterministic_algorithms_enabled, + disable_flex_compile, + disable_flex_compile_backprop, + get_memory_usage_preference, + is_flex_compile_allowed, + is_flex_compile_backprop_allowed, + is_kv_parallelism_in_fused_na_enabled, + is_memory_usage_default, + is_memory_usage_strict, + is_memory_usage_unrestricted, + set_memory_usage_preference, + use_deterministic_algorithms, + use_kv_parallelism_in_fused_na, +) +from .functional import attention, merge_attentions, na1d, na2d, na3d +from .modules import ( + NeighborhoodAttention1D, + NeighborhoodAttention2D, + NeighborhoodAttention3D, +) +from .version import __version__ + +__all__ = [ + "__version__", + "NeighborhoodAttention1D", + "NeighborhoodAttention2D", + "NeighborhoodAttention3D", + "are_deterministic_algorithms_enabled", + "use_deterministic_algorithms", + "use_kv_parallelism_in_fused_na", + "is_kv_parallelism_in_fused_na_enabled", + "set_memory_usage_preference", + "get_memory_usage_preference", + "is_memory_usage_default", + "is_memory_usage_strict", + "is_memory_usage_unrestricted", + "is_flex_compile_allowed", + "is_flex_compile_backprop_allowed", + "allow_flex_compile", + "allow_flex_compile_backprop", + "disable_flex_compile", + "disable_flex_compile_backprop", + "get_bwd_configs_for_cutlass_fmha", + "get_bwd_configs_for_cutlass_fna", + "get_configs_for_cutlass_fmha", + "get_configs_for_cutlass_fna", + "get_configs_for_cutlass_hopper_fmha", + "get_bwd_configs_for_cutlass_hopper_fmha", + "get_configs_for_cutlass_hopper_fna", + "get_bwd_configs_for_cutlass_hopper_fna", + "get_bwd_configs_for_cutlass_blackwell_fmha", + "get_bwd_configs_for_cutlass_blackwell_fna", + "get_configs_for_cutlass_blackwell_fmha", + "get_configs_for_cutlass_blackwell_fna", + "get_configs_for_flex_fmha", + "get_configs_for_flex_fna", + "HAS_LIBNATTEN", + "na1d", + "na2d", + "na3d", + "attention", + "merge_attentions", +] diff --git a/natten/torch-ext/natten/_environment.py b/natten/torch-ext/natten/_environment.py new file mode 100644 index 00000000..caae85bb --- /dev/null +++ b/natten/torch-ext/natten/_environment.py @@ -0,0 +1,59 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +from ._libnatten import HAS_LIBNATTEN # noqa: F401 +from .utils.environment import ( + _IS_CUDA_AVAILABLE, + _IS_TORCH_COMPILE_SUPPORTED, + _TORCH_VERSION, + parse_env_flag, + parse_env_int, + parse_env_str, +) + +# Default tokperm implementation; choices: +# NATTEN_TOKPERM_DEFAULT_IMPL="cutlass" +# NATTEN_TOKPERM_DEFAULT_IMPL="torch" +USE_TORCH_IMPL_DEFAULT = ( + parse_env_str("NATTEN_TOKPERM_DEFAULT_IMPL", "cutlass") == "torch" +) + +# Unit tests +_RUN_EXTENDED_TESTS = parse_env_flag("NATTEN_RUN_EXTENDED_TESTS", False) +_RUN_FLEX_TESTS = parse_env_flag("NATTEN_RUN_FLEX_TESTS", True) +_NUM_RAND_SWEEP_TESTS = parse_env_int("NATTEN_RAND_SWEEP_TESTS", 1000) + +# Profiler +DISABLE_TQDM = parse_env_flag("NATTEN_DISABLE_TQDM", False) + + +__all__ = [ + "HAS_LIBNATTEN", + "_IS_CUDA_AVAILABLE", + "_IS_TORCH_COMPILE_SUPPORTED", + "DISABLE_TQDM", + "_RUN_FLEX_TESTS", + "_RUN_FLEX_TESTS", + "_NUM_RAND_SWEEP_TESTS", + "_TORCH_VERSION", +] diff --git a/natten/torch-ext/natten/_libnatten/__init__.py b/natten/torch-ext/natten/_libnatten/__init__.py new file mode 100644 index 00000000..9566f8d4 --- /dev/null +++ b/natten/torch-ext/natten/_libnatten/__init__.py @@ -0,0 +1,109 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# +# kernel-builder port: libnatten is always compiled into this kernel; the +# upstream stub fallback path is not needed. + +import torch # noqa: F401 + +HAS_LIBNATTEN = True + +from .torch_wrappers import ( + blackwell_fmha_backward, + blackwell_fmha_forward, + blackwell_na1d_backward, + blackwell_na1d_forward, + blackwell_na2d_backward, + blackwell_na2d_forward, + blackwell_na3d_backward, + blackwell_na3d_forward, + compute_delta, + fmha_backward, + fmha_forward, + hopper_fmha_backward, + hopper_fmha_forward, + hopper_na1d_backward, + hopper_na1d_forward, + hopper_na2d_backward, + hopper_na2d_forward, + hopper_na3d_backward, + hopper_na3d_forward, + na1d_backward, + na1d_forward, + na2d_backward, + na2d_forward, + na3d_backward, + na3d_forward, + reference_na1d_backward, + reference_na1d_forward, + reference_na2d_backward, + reference_na2d_forward, + reference_na3d_backward, + reference_na3d_forward, + token_permute_1d, + token_permute_2d, + token_permute_3d, + token_unpermute_1d, + token_unpermute_2d, + token_unpermute_3d, +) + +__all__ = [ + "HAS_LIBNATTEN", + "blackwell_fmha_backward", + "blackwell_fmha_forward", + "blackwell_na1d_backward", + "blackwell_na1d_forward", + "blackwell_na2d_backward", + "blackwell_na2d_forward", + "blackwell_na3d_backward", + "blackwell_na3d_forward", + "compute_delta", + "fmha_backward", + "fmha_forward", + "hopper_fmha_backward", + "hopper_fmha_forward", + "hopper_na1d_backward", + "hopper_na1d_forward", + "hopper_na2d_backward", + "hopper_na2d_forward", + "hopper_na3d_backward", + "hopper_na3d_forward", + "na1d_backward", + "na1d_forward", + "na2d_backward", + "na2d_forward", + "na3d_backward", + "na3d_forward", + "reference_na1d_backward", + "reference_na1d_forward", + "reference_na2d_backward", + "reference_na2d_forward", + "reference_na3d_backward", + "reference_na3d_forward", + "token_permute_1d", + "token_permute_2d", + "token_permute_3d", + "token_unpermute_1d", + "token_unpermute_2d", + "token_unpermute_3d", +] diff --git a/natten/torch-ext/natten/_libnatten/torch_wrappers.py b/natten/torch-ext/natten/_libnatten/torch_wrappers.py new file mode 100644 index 00000000..16e5b62b --- /dev/null +++ b/natten/torch-ext/natten/_libnatten/torch_wrappers.py @@ -0,0 +1,1006 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# +# kernel-builder port of upstream `natten/_libnatten/torch_wrappers.py`. +# +# Upstream registers Python `torch.library.custom_op`s that allocate outputs +# and call into the pybind11 `libnatten` extension. In this port the ops are +# registered in C++ (`torch-ext/torch_binding.cpp`) as out-variant ops under +# the build-time namespace exposed through `.._ops`. The functions here keep +# the exact upstream calling conventions (allocate outputs, handle kv-split +# defaults, varlen zero-init) and call the C++ ops, and each C++ op gets a +# fake (meta) registration so the whole surface stays torch.compile-safe. +# +# Schema conventions of the C++ ops: +# - `kernel_size`/`stride`/`dilation`/tile shapes are `int[]`. +# - Multi-dimensional causal masks are passed as `int[]` (0/1) because +# boolean arrays are less uniformly supported in op schemas. +# - `scale` is a `float`. + +import math +from typing import Optional, Sequence, Tuple + +import torch +from torch import Tensor + +from .._ops import add_op_namespace_prefix, ops +from ..utils.tuples import ceil_div_tuple, mul_tuple + +register_fake = torch.library.register_fake + + +def maybe_contiguous(x): + return x.contiguous() + + +def _ints(v: Sequence) -> list: + return [int(x) for x in v] + + +################################################################################ +############################ Fake (meta) registration ########################## +################################################################################ +# All C++ ops are out-variant: they only mutate output arguments and return +# nothing, so their fake impls are no-ops. Shape inference happens in the +# Python wrappers below, which allocate the outputs. + + +def _register_noop_fake(op_name: str) -> None: + def _fake(*args, **kwargs) -> None: + return None + + register_fake(add_op_namespace_prefix(op_name))(_fake) + + +for _na_dim in (1, 2, 3): + for _prefix in ("", "hopper_", "blackwell_", "reference_"): + _register_noop_fake(f"{_prefix}na{_na_dim}d_forward") + _register_noop_fake(f"{_prefix}na{_na_dim}d_backward") + _register_noop_fake(f"token_permute_{_na_dim}d") + _register_noop_fake(f"token_unpermute_{_na_dim}d") + +for _prefix in ("", "hopper_", "blackwell_"): + _register_noop_fake(f"{_prefix}fmha_forward") + _register_noop_fake(f"{_prefix}fmha_backward") + +_register_noop_fake("compute_delta") + + +################################################################################ +################################### FMHA ops ################################### +################################################################################ + + +def blackwell_fmha_forward( + query: Tensor, + key: Tensor, + value: Tensor, + is_causal: bool, + scale: float, + q_tile_size: int, + kv_tile_size: int, + run_persistent_kernel: bool, + cumulative_seqlen_Q: Optional[Tensor], + cumulative_seqlen_KV: Optional[Tensor], + max_seqlen_Q: int, + max_seqlen_KV: int, +) -> Tuple[Tensor, Tensor]: + query, key, value = [maybe_contiguous(x) for x in (query, key, value)] + + output_shape = [s for s in query.shape[:-1]] + [value.shape[-1]] + + # NOTE: always zero-init outputs when doing varlen for safety + is_varlen = cumulative_seqlen_Q is not None + init_fn = torch.zeros if is_varlen else torch.empty + + output = init_fn(output_shape, device=query.device, dtype=query.dtype) + logsumexp = init_fn(query.shape[:-1], dtype=torch.float32, device=query.device) + + # Skip kernel launch when all sequences are empty + if is_varlen and max_seqlen_Q == 0 and max_seqlen_KV == 0: + return output, logsumexp + + ops.blackwell_fmha_forward( + output, + query, + key, + value, + logsumexp, + bool(is_causal), + float(scale), + int(q_tile_size), + int(kv_tile_size), + bool(run_persistent_kernel), + cumulative_seqlen_Q, + cumulative_seqlen_KV, + int(max_seqlen_Q), + int(max_seqlen_KV), + ) + + return output, logsumexp + + +def blackwell_fmha_backward( + query: Tensor, + key: Tensor, + value: Tensor, + output: Tensor, + d_output: Tensor, + logsumexp: Tensor, + is_causal: bool, + scale: float, + q_tile_size: int, + kv_tile_size: int, + cumulative_seqlen_Q: Optional[Tensor], + cumulative_seqlen_KV: Optional[Tensor], + max_seqlen_Q: int, + max_seqlen_KV: int, + deterministic: bool, +) -> Tuple[Tensor, Tensor, Tensor]: + query, key, value = [maybe_contiguous(x) for x in (query, key, value)] + output, d_output, logsumexp = [ + maybe_contiguous(x) for x in (output, d_output, logsumexp) + ] + + # NOTE: always zero-init outputs when doing varlen for safety + is_varlen = cumulative_seqlen_Q is not None + init_fn = torch.zeros_like if is_varlen else torch.empty_like + + d_query = init_fn(query) + d_key = init_fn(key) + d_value = init_fn(value) + + # Skip kernel launch when all sequences are empty + if is_varlen and max_seqlen_Q == 0 and max_seqlen_KV == 0: + return d_query, d_key, d_value + + ops.blackwell_fmha_backward( + d_query, + d_key, + d_value, + query, + key, + value, + output, + d_output, + logsumexp, + bool(is_causal), + float(scale), + int(q_tile_size), + int(kv_tile_size), + cumulative_seqlen_Q, + cumulative_seqlen_KV, + int(max_seqlen_Q), + int(max_seqlen_KV), + bool(deterministic), + ) + + return d_query, d_key, d_value + + +def hopper_fmha_forward( + query: Tensor, + key: Tensor, + value: Tensor, + is_causal: bool, + scale: float, + q_tile_size: int, + kv_tile_size: int, + kernel_schedule_int: int, + cumulative_seqlen_Q: Optional[Tensor], + cumulative_seqlen_KV: Optional[Tensor], + max_seqlen_Q: int, + max_seqlen_KV: int, +) -> Tuple[Tensor, Tensor]: + query, key, value = [maybe_contiguous(x) for x in (query, key, value)] + + output_shape = [s for s in query.shape[:-1]] + [value.shape[-1]] + + # NOTE: always zero-init outputs when doing varlen for safety + is_varlen = cumulative_seqlen_Q is not None + init_fn = torch.zeros if is_varlen else torch.empty + + output = init_fn(output_shape, device=query.device, dtype=query.dtype) + logsumexp = init_fn(query.shape[:-1], dtype=torch.float32, device=query.device) + + # Skip kernel launch when all sequences are empty + if is_varlen and max_seqlen_Q == 0 and max_seqlen_KV == 0: + return output, logsumexp + + ops.hopper_fmha_forward( + output, + query, + key, + value, + logsumexp, + bool(is_causal), + float(scale), + int(q_tile_size), + int(kv_tile_size), + int(kernel_schedule_int), + cumulative_seqlen_Q, + cumulative_seqlen_KV, + int(max_seqlen_Q), + int(max_seqlen_KV), + ) + + return output, logsumexp + + +def hopper_fmha_backward( + query: Tensor, + key: Tensor, + value: Tensor, + output: Tensor, + d_output: Tensor, + logsumexp: Tensor, + is_causal: bool, + scale: float, + q_tile_size: int, + kv_tile_size: int, + cumulative_seqlen_Q: Optional[Tensor], + cumulative_seqlen_KV: Optional[Tensor], + max_seqlen_Q: int, + max_seqlen_KV: int, +) -> Tuple[Tensor, Tensor, Tensor]: + query, key, value = [maybe_contiguous(x) for x in (query, key, value)] + output, d_output, logsumexp = [ + maybe_contiguous(x) for x in (output, d_output, logsumexp) + ] + + # NOTE: always zero-init outputs when doing varlen for safety + is_varlen = cumulative_seqlen_Q is not None + init_fn = torch.zeros_like if is_varlen else torch.empty_like + + d_query = init_fn(query) + d_key = init_fn(key) + d_value = init_fn(value) + + # Skip kernel launch when all sequences are empty + if is_varlen and max_seqlen_Q == 0 and max_seqlen_KV == 0: + return d_query, d_key, d_value + + ops.hopper_fmha_backward( + d_query, + d_key, + d_value, + query, + key, + value, + output, + d_output, + logsumexp, + bool(is_causal), + float(scale), + int(q_tile_size), + int(kv_tile_size), + cumulative_seqlen_Q, + cumulative_seqlen_KV, + int(max_seqlen_Q), + int(max_seqlen_KV), + ) + + return d_query, d_key, d_value + + +def fmha_forward( + query: Tensor, + key: Tensor, + value: Tensor, + is_causal: bool, + scale: float, + q_tile_size: int, + kv_tile_size: int, + cumulative_seqlen_Q: Optional[Tensor], + cumulative_seqlen_KV: Optional[Tensor], + max_seqlen_Q: int, + max_seqlen_KV: int, +) -> Tuple[Tensor, Tensor]: + query, key, value = [maybe_contiguous(x) for x in (query, key, value)] + + output_shape = [s for s in query.shape[:-1]] + [value.shape[-1]] + + # NOTE: always zero-init outputs when doing varlen for safety + is_varlen = cumulative_seqlen_Q is not None + init_fn = torch.zeros if is_varlen else torch.empty + + output = init_fn(output_shape, device=query.device, dtype=query.dtype) + logsumexp = init_fn(query.shape[:-1], dtype=torch.float32, device=query.device) + + # Skip kernel launch when all sequences are empty + if is_varlen and max_seqlen_Q == 0 and max_seqlen_KV == 0: + return output, logsumexp + + ops.fmha_forward( + output, + query, + key, + value, + logsumexp, + bool(is_causal), + float(scale), + int(q_tile_size), + int(kv_tile_size), + cumulative_seqlen_Q, + cumulative_seqlen_KV, + int(max_seqlen_Q), + int(max_seqlen_KV), + ) + + return output, logsumexp + + +def fmha_backward( + query: Tensor, + key: Tensor, + value: Tensor, + output: Tensor, + d_output: Tensor, + logsumexp: Tensor, + is_causal: bool, + scale: float, + q_tile_size: int, + kv_tile_size: int, + num_kv_splits: Optional[int], + compute_delta_with_pt: bool, + cumulative_seqlen_Q: Optional[Tensor], + cumulative_seqlen_KV: Optional[Tensor], + max_seqlen_Q: int, + max_seqlen_KV: int, + deterministic: bool, +) -> Tuple[Tensor, Tensor, Tensor]: + from ..backends.configs.cutlass.backward_knobs import check_fmha_kv_splits + + query, key, value = [maybe_contiguous(x) for x in (query, key, value)] + output, d_output, logsumexp = [ + maybe_contiguous(x) for x in (output, d_output, logsumexp) + ] + + # NOTE: always zero-init outputs when doing varlen for safety + is_varlen = cumulative_seqlen_Q is not None + init_fn = torch.zeros_like if is_varlen else torch.empty_like + + d_query = init_fn(query) + d_key = init_fn(key) + d_value = init_fn(value) + + # Skip kernel launch when all sequences are empty + if is_varlen and max_seqlen_Q == 0 and max_seqlen_KV == 0: + return d_query, d_key, d_value + + if deterministic: + # Torch reduction seems to have slight reproducibility issues, even with determinism on + compute_delta_with_pt = False + # TODO: this is the only way to get determinism in this kernel, but it's very slow + num_kv_splits = 1 + else: + # Compute default kv_splits if not specified + # max_seqlen must be at least 2 to satisfy static checks that are just too complicated to + # relax at this point. Kernel launch will be skipped if max_seqlen is 0 anyway. Prior checks + # should prevent negative max seqlens. + max_seqlen = max(2, max_seqlen_KV) if is_varlen else None + num_kv_splits = check_fmha_kv_splits( + kv_splits=num_kv_splits, + input_tensor=key, + kv_tile_size=kv_tile_size, + deterministic=deterministic, + max_seqlen=max_seqlen, + ) + + ops.fmha_backward( + d_query, + d_key, + d_value, + query, + key, + value, + output, + d_output, + logsumexp, + bool(is_causal), + float(scale), + int(q_tile_size), + int(kv_tile_size), + int(num_kv_splits), + bool(compute_delta_with_pt), + cumulative_seqlen_Q, + cumulative_seqlen_KV, + int(max_seqlen_Q), + int(max_seqlen_KV), + ) + + return d_query, d_key, d_value + + +################################################################################ +################################### FNA ops ################################### +################################################################################ + + +def make_blackwell_fna_ops(na_dim): + fwd_op = getattr(ops, f"blackwell_na{na_dim}d_forward") + bwd_op = getattr(ops, f"blackwell_na{na_dim}d_backward") + + def blackwell_fna_forward( + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size, + stride, + dilation, + is_causal, + scale: float, + q_shape, + kv_shape, + qkv_shape, + q_tile_shape, + kv_tile_shape, + run_persistent_kernel: bool, + ) -> Tuple[Tensor, Tensor]: + query, key, value = [maybe_contiguous(x) for x in (query, key, value)] + + output_shape = [s for s in query.shape[:-1]] + [value.shape[-1]] + output = torch.empty(output_shape, device=query.device, dtype=query.dtype) + + logsumexp = torch.empty( + query.shape[:-1], dtype=torch.float32, device=query.device + ) + + fwd_op( + output, + query, + key, + value, + logsumexp, + _ints(kernel_size), + _ints(stride), + _ints(dilation), + _ints(is_causal), + float(scale), + _ints(q_shape), + _ints(kv_shape), + _ints(qkv_shape), + _ints(q_tile_shape), + _ints(kv_tile_shape), + bool(run_persistent_kernel), + ) + + return output, logsumexp + + def blackwell_fna_backward( + query: Tensor, + key: Tensor, + value: Tensor, + output: Tensor, + d_output: Tensor, + logsumexp: Tensor, + kernel_size, + stride, + dilation, + is_causal, + scale: float, + q_shape, + kv_shape, + qkv_shape, + q_tile_shape, + kv_tile_shape, + ) -> Tuple[Tensor, Tensor, Tensor]: + query, key, value = [maybe_contiguous(x) for x in (query, key, value)] + output, d_output, logsumexp = [ + maybe_contiguous(x) for x in (output, d_output, logsumexp) + ] + + d_query = torch.empty_like(query) + d_key = torch.empty_like(key) + d_value = torch.empty_like(value) + + bwd_op( + d_query, + d_key, + d_value, + query, + key, + value, + output, + d_output, + logsumexp, + _ints(kernel_size), + _ints(stride), + _ints(dilation), + _ints(is_causal), + float(scale), + _ints(q_shape), + _ints(kv_shape), + _ints(qkv_shape), + _ints(q_tile_shape), + _ints(kv_tile_shape), + ) + + return d_query, d_key, d_value + + return blackwell_fna_forward, blackwell_fna_backward + + +def make_hopper_fna_ops(na_dim): + fwd_op = getattr(ops, f"hopper_na{na_dim}d_forward") + bwd_op = getattr(ops, f"hopper_na{na_dim}d_backward") + + def hopper_fna_forward( + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size, + stride, + dilation, + is_causal, + scale: float, + q_shape, + kv_shape, + qkv_shape, + q_tile_shape, + kv_tile_shape, + kernel_schedule_int: int, + ) -> Tuple[Tensor, Tensor]: + query, key, value = [maybe_contiguous(x) for x in (query, key, value)] + + output_shape = [s for s in query.shape[:-1]] + [value.shape[-1]] + output = torch.empty(output_shape, device=query.device, dtype=query.dtype) + + logsumexp = torch.empty( + query.shape[:-1], dtype=torch.float32, device=query.device + ) + + fwd_op( + output, + query, + key, + value, + logsumexp, + _ints(kernel_size), + _ints(stride), + _ints(dilation), + _ints(is_causal), + float(scale), + _ints(q_shape), + _ints(kv_shape), + _ints(qkv_shape), + _ints(q_tile_shape), + _ints(kv_tile_shape), + int(kernel_schedule_int), + ) + + return output, logsumexp + + def hopper_fna_backward( + query: Tensor, + key: Tensor, + value: Tensor, + output: Tensor, + d_output: Tensor, + logsumexp: Tensor, + kernel_size, + stride, + dilation, + is_causal, + scale: float, + q_shape, + kv_shape, + qkv_shape, + q_tile_shape, + kv_tile_shape, + ) -> Tuple[Tensor, Tensor, Tensor]: + query, key, value = [maybe_contiguous(x) for x in (query, key, value)] + output, d_output, logsumexp = [ + maybe_contiguous(x) for x in (output, d_output, logsumexp) + ] + + d_query = torch.empty_like(query) + d_key = torch.empty_like(key) + d_value = torch.empty_like(value) + + bwd_op( + d_query, + d_key, + d_value, + query, + key, + value, + output, + d_output, + logsumexp, + _ints(kernel_size), + _ints(stride), + _ints(dilation), + _ints(is_causal), + float(scale), + _ints(q_shape), + _ints(kv_shape), + _ints(qkv_shape), + _ints(q_tile_shape), + _ints(kv_tile_shape), + ) + + return d_query, d_key, d_value + + return hopper_fna_forward, hopper_fna_backward + + +def make_fna_ops(na_dim): + fwd_op = getattr(ops, f"na{na_dim}d_forward") + bwd_op = getattr(ops, f"na{na_dim}d_backward") + + def fna_forward( + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size, + stride, + dilation, + is_causal, + scale: float, + q_tile_shape, + kv_tile_shape, + ) -> Tuple[Tensor, Tensor]: + query, key, value = [maybe_contiguous(x) for x in (query, key, value)] + + output_shape = [s for s in query.shape[:-1]] + [value.shape[-1]] + output = torch.empty(output_shape, device=query.device, dtype=query.dtype) + + logsumexp = torch.empty( + query.shape[:-1], dtype=torch.float32, device=query.device + ) + + fwd_op( + output, + query, + key, + value, + logsumexp, + _ints(kernel_size), + _ints(stride), + _ints(dilation), + _ints(is_causal), + float(scale), + _ints(q_tile_shape), + _ints(kv_tile_shape), + ) + + return output, logsumexp + + def fna_backward( + query: Tensor, + key: Tensor, + value: Tensor, + output: Tensor, + d_output: Tensor, + logsumexp: Tensor, + kernel_size, + stride, + dilation, + is_causal, + scale: float, + q_tile_shape, + kv_tile_shape, + num_kv_splits, + compute_delta_with_pt: bool, + deterministic: bool, + ) -> Tuple[Tensor, Tensor, Tensor]: + from ..backends.configs.cutlass.backward_knobs import check_fna_kv_splits + + query, key, value = [maybe_contiguous(x) for x in (query, key, value)] + output, d_output, logsumexp = [ + maybe_contiguous(x) for x in (output, d_output, logsumexp) + ] + + d_query = torch.empty_like(query) + d_key = torch.empty_like(key) + d_value = torch.empty_like(value) + + if deterministic: + # Torch reduction seems to have slight reproducibility issues, even with determinism on + compute_delta_with_pt = False + # TODO: this is the only way to get determinism in this kernel, but it's very slow + num_kv_splits = tuple(1 for _ in range(na_dim)) + else: + # Compute default kv_splits if not specified + num_kv_splits = check_fna_kv_splits( + kv_splits=tuple(num_kv_splits) if num_kv_splits is not None else None, + input_tensor=key, + kv_tile_shape=tuple(kv_tile_shape), + deterministic=deterministic, + dilation=tuple(dilation), + ) + + bwd_op( + d_query, + d_key, + d_value, + query, + key, + value, + output, + d_output, + logsumexp, + _ints(kernel_size), + _ints(stride), + _ints(dilation), + _ints(is_causal), + float(scale), + _ints(q_tile_shape), + _ints(kv_tile_shape), + _ints(num_kv_splits), + bool(compute_delta_with_pt), + ) + + return d_query, d_key, d_value + + return fna_forward, fna_backward + + +def make_reference_fna_ops(na_dim): + fwd_op = getattr(ops, f"reference_na{na_dim}d_forward") + bwd_op = getattr(ops, f"reference_na{na_dim}d_backward") + + def reference_fna_forward( + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size, + stride, + dilation, + is_causal, + scale: float, + qkv_shape, + num_extra_kv: int, + ) -> Tuple[Tensor, Tensor]: + query, key, value = [maybe_contiguous(x) for x in (query, key, value)] + + output_shape = [s for s in query.shape[:-1]] + [value.shape[-1]] + output = torch.empty(output_shape, device=query.device, dtype=query.dtype) + + logsumexp = torch.empty( + query.shape[:-1], dtype=torch.float32, device=query.device + ) + + fwd_op( + output, + query, + key, + value, + logsumexp, + _ints(kernel_size), + _ints(stride), + _ints(dilation), + _ints(is_causal), + float(scale), + _ints(qkv_shape), + int(num_extra_kv), + ) + + return output, logsumexp + + def reference_fna_backward( + query: Tensor, + key: Tensor, + value: Tensor, + output: Tensor, + d_output: Tensor, + logsumexp: Tensor, + kernel_size, + stride, + dilation, + is_causal, + scale: float, + qkv_shape, + num_extra_kv: int, + ) -> Tuple[Tensor, Tensor, Tensor]: + query, key, value = [maybe_contiguous(x) for x in (query, key, value)] + output, d_output, logsumexp = [ + maybe_contiguous(x) for x in (output, d_output, logsumexp) + ] + + d_query = torch.empty_like(query) + d_key = torch.empty_like(key) + d_value = torch.empty_like(value) + + bwd_op( + d_query, + d_key, + d_value, + query, + key, + value, + output, + d_output, + logsumexp, + _ints(kernel_size), + _ints(stride), + _ints(dilation), + _ints(is_causal), + float(scale), + _ints(qkv_shape), + int(num_extra_kv), + ) + + return d_query, d_key, d_value + + return reference_fna_forward, reference_fna_backward + + +################################################################################ +################################# TokPerm ops ################################# +################################################################################ + + +def make_token_permute_ops(na_dim): + permute_op = getattr(ops, f"token_permute_{na_dim}d") + unpermute_op = getattr(ops, f"token_unpermute_{na_dim}d") + + def token_permute( + input_tensor: Tensor, + tile_shape, + dilation, + flip_tiled_dims: bool, + ) -> Tensor: + input_tensor = maybe_contiguous(input_tensor) + + token_layout = tuple(x for x in input_tensor.shape[1 : na_dim + 1]) + token_layout_padded = mul_tuple( + mul_tuple( + ceil_div_tuple(ceil_div_tuple(token_layout, tile_shape), dilation), + dilation, + ), + tile_shape, + ) + output_shape = [ + input_tensor.shape[0], + math.prod(token_layout_padded), + input_tensor.shape[-2], + input_tensor.shape[-1], + ] + output = torch.empty( + output_shape, device=input_tensor.device, dtype=input_tensor.dtype + ) + permute_op( + output, + input_tensor, + _ints(tile_shape), + _ints(dilation), + bool(flip_tiled_dims), + ) + + # Fold dilation in batch dimension so that attention is correct. + output = output.reshape( + input_tensor.shape[0] * math.prod(dilation), + -1, + input_tensor.shape[-2], + input_tensor.shape[-1], + ) + + return output + + def token_unpermute( + input_tensor: Tensor, + token_layout_shape, + tile_shape, + dilation, + flip_tiled_dims: bool, + ) -> Tensor: + input_tensor = maybe_contiguous(input_tensor) + + # Unfold dilation in batch dimension + num_dilation_groups = math.prod(dilation) + assert input_tensor.shape[0] % num_dilation_groups == 0 + input_tensor = input_tensor.reshape( + input_tensor.shape[0] // num_dilation_groups, + -1, + input_tensor.shape[-2], + input_tensor.shape[-1], + ) + + output_shape = [ + input_tensor.shape[0], + *token_layout_shape, + input_tensor.shape[-2], + input_tensor.shape[-1], + ] + output = torch.empty( + output_shape, device=input_tensor.device, dtype=input_tensor.dtype + ) + unpermute_op( + output, + input_tensor, + _ints(tile_shape), + _ints(dilation), + bool(flip_tiled_dims), + ) + + return output + + return token_permute, token_unpermute + + +(blackwell_na1d_forward, blackwell_na1d_backward) = make_blackwell_fna_ops(1) +(blackwell_na2d_forward, blackwell_na2d_backward) = make_blackwell_fna_ops(2) +(blackwell_na3d_forward, blackwell_na3d_backward) = make_blackwell_fna_ops(3) + +(hopper_na1d_forward, hopper_na1d_backward) = make_hopper_fna_ops(1) +(hopper_na2d_forward, hopper_na2d_backward) = make_hopper_fna_ops(2) +(hopper_na3d_forward, hopper_na3d_backward) = make_hopper_fna_ops(3) + +(na1d_forward, na1d_backward) = make_fna_ops(1) +(na2d_forward, na2d_backward) = make_fna_ops(2) +(na3d_forward, na3d_backward) = make_fna_ops(3) + +(reference_na1d_forward, reference_na1d_backward) = make_reference_fna_ops(1) +(reference_na2d_forward, reference_na2d_backward) = make_reference_fna_ops(2) +(reference_na3d_forward, reference_na3d_backward) = make_reference_fna_ops(3) + +(token_permute_1d, token_unpermute_1d) = make_token_permute_ops(1) +(token_permute_2d, token_unpermute_2d) = make_token_permute_ops(2) +(token_permute_3d, token_unpermute_3d) = make_token_permute_ops(3) + + +# This is only used in unit tests, and not even auto-diffable +def compute_delta(out: Tensor, d_out: Tensor, delta: Tensor) -> None: + ops.compute_delta(out, d_out, delta) + + +__all__ = [ + "blackwell_fmha_backward", + "blackwell_fmha_forward", + "blackwell_na1d_backward", + "blackwell_na1d_forward", + "blackwell_na2d_backward", + "blackwell_na2d_forward", + "blackwell_na3d_backward", + "blackwell_na3d_forward", + "compute_delta", + "fmha_backward", + "fmha_forward", + "hopper_fmha_backward", + "hopper_fmha_forward", + "hopper_na1d_backward", + "hopper_na1d_forward", + "hopper_na2d_backward", + "hopper_na2d_forward", + "hopper_na3d_backward", + "hopper_na3d_forward", + "na1d_backward", + "na1d_forward", + "na2d_backward", + "na2d_forward", + "na3d_backward", + "na3d_forward", + "reference_na1d_backward", + "reference_na1d_forward", + "reference_na2d_backward", + "reference_na2d_forward", + "reference_na3d_backward", + "reference_na3d_forward", + "token_permute_1d", + "token_permute_2d", + "token_permute_3d", + "token_unpermute_1d", + "token_unpermute_2d", + "token_unpermute_3d", +] diff --git a/natten/torch-ext/natten/attn_merge.py b/natten/torch-ext/natten/attn_merge.py new file mode 100644 index 00000000..e71cc799 --- /dev/null +++ b/natten/torch-ext/natten/attn_merge.py @@ -0,0 +1,292 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +import functools +from typing import List, Tuple + +import torch +from torch import Tensor +from torch.amp import custom_bwd, custom_fwd +from torch.autograd import Function + +amp_fwd = functools.partial(custom_fwd, device_type="cuda") +amp_bwd = functools.partial(custom_bwd, device_type="cuda") + +from ._environment import _IS_TORCH_COMPILE_SUPPORTED + + +def _maybe_torch_compile(*args, **kwargs): + def decorator(f): + if _IS_TORCH_COMPILE_SUPPORTED: + return torch.compile(f, *args, **kwargs) + return f + + return decorator + + +# TODO: if use cases for this grow, we might want to do a custom kernel +def _merge_attentions_fn( + outputs: List[Tensor], lse_tensors: List[Tensor] +) -> Tuple[Tensor, Tensor]: + + assert len(outputs) >= 2, "Expected at least two tensors." + num_splits = len(outputs) + assert ( + len(lse_tensors) == num_splits + ), "Expected number of outputs and LSE tensors to match." + + assert all( + output.dim() == 4 and output.is_contiguous() for output in outputs + ), "Output tensors must be rank-4 tensors with (batch, seq, heads, dim) contiguous layout." + + batch, seqlen, heads, dim = outputs[0].shape + + assert all( + [x for x in output.shape] == [batch, seqlen, heads, dim] for output in outputs + ), "Output tensors must match in shape." + + assert all( + lse.dim() == 3 + and lse.is_contiguous() + and [x for x in lse.shape] == [batch, seqlen, heads] + for lse in lse_tensors + ), "LSE tensors must be rank-3 tensors with (batch, seq, heads) contiguous layout, and match in shape." + + accum_type = torch.float32 + output_type = outputs[0].dtype + + lse_tensors = [lse.to(accum_type).unsqueeze(-1) for lse in lse_tensors] + + outputs = [output.to(accum_type) for output in outputs] + + # New approach based on https://github.com/zhuzilin/ring-flash-attention/pull/34 + output = outputs[0] - torch.nn.functional.sigmoid( + lse_tensors[1] - lse_tensors[0] + ) * (outputs[0] - outputs[1]) + logsumexp = lse_tensors[0] - torch.nn.functional.logsigmoid( + lse_tensors[0] - lse_tensors[1] + ) + for i in range(2, num_splits): + output = output - torch.nn.functional.sigmoid(lse_tensors[i] - logsumexp) * ( + output - outputs[i] + ) + logsumexp = logsumexp - torch.nn.functional.logsigmoid( + logsumexp - lse_tensors[i] + ) + + output = output.to(output_type) + logsumexp = logsumexp.squeeze(-1) + + assert logsumexp.dim() == 3 + assert logsumexp.shape[0] == batch + assert logsumexp.shape[1] == seqlen + assert logsumexp.shape[2] == heads + + return output, logsumexp + + +@_maybe_torch_compile(fullgraph=True) +def _merge_attentions_compile( + outputs: List[Tensor], lse_tensors: List[Tensor] +) -> Tuple[Tensor, Tensor]: + return _merge_attentions_fn(outputs, lse_tensors) + + +def _merge_attentions_op( + outputs: List[Tensor], lse_tensors: List[Tensor], torch_compile: bool = True +) -> Tuple[Tensor, Tensor]: + + if not torch_compile: + return _merge_attentions_fn( + [output.contiguous() for output in outputs], + [lse.contiguous() for lse in lse_tensors], + ) + + return _merge_attentions_compile( + [output.contiguous() for output in outputs], + [lse.contiguous() for lse in lse_tensors], + ) + + +class MergeAttentionsAutogradFn(Function): + @staticmethod + @amp_fwd + def forward( + ctx, + *args, + ) -> Tuple[Tensor, Tensor]: + + assert len(args) >= 5, ( + "Expected at least 5 args (two outputs, two lse tensors, 1 torch compile flag) " + + f"in attention merge, got {len(args)}." + ) + assert (len(args) - 1) % 2 == 0, ( + f"Expected pairs of outputs and lse tensors, got {len(args)-1} args " + + "(excluding torch compile flag)" + ) + num_pairs = (len(args) - 1) // 2 + assert num_pairs >= 2 + + torch_compile = args[-1] + outputs = args[:num_pairs] + lses = args[num_pairs:-1] + + assert len(outputs) == len(lses), ( + "Expected the same number of outputs as logsumexp tensors, " + + f"got {len(outputs)=}, {len(lses)=}" + ) + + merged_output, merged_lse = _merge_attentions_op( + outputs, # type: ignore[arg-type] + lses, # type: ignore[arg-type] + torch_compile=torch_compile, + ) + + ctx.num_pairs = num_pairs + ctx.save_for_backward(merged_output, merged_lse, *outputs, *lses) + + return merged_output, merged_lse + + @staticmethod + @amp_bwd + def backward(ctx, grad_out: Tensor, grad_lse: Tensor) -> Tuple: + + num_pairs = ctx.num_pairs + merged_output, merged_lse = ctx.saved_tensors[:2] + outputs = ctx.saved_tensors[2 : num_pairs + 2] + lses = ctx.saved_tensors[num_pairs + 2 :] + + # Outputs and LSEs from the originating attention ops must be replaced with + # the merged ones inplace so that we get correct behavior, and not break torch.compile + # graphs in the process. + for output, lse in zip(outputs, lses): + output.data.copy_(merged_output.data.reshape(output.shape)) + lse.data.copy_(merged_lse.data.reshape(lse.shape)) + + return ( + *(grad_out for _ in range(num_pairs)), + *(grad_lse for _ in range(num_pairs)), + None, + ) + + +def merge_attentions( + outputs: List[Tensor], + lse_tensors: List[Tensor], + torch_compile: bool = True, + use_autograd_fix: bool = True, +) -> Tuple[Tensor, Tensor]: + """Takes multiple attention *outputs* originating from the same query tensor, and their + corresponding logsumexps, and merges them as if their context (key/value pair) had been + concatenated. + + This operation is used to implement cross-neighborhood attention, and can also be used for + distributed setups, such as context-parallelism. + + This operation also attempts to use `torch.compile` to fuse the elementwise operations. This + can be disabled by passing `torch_compile=False`. + + Parameters: + outputs (List[Tensor]): List of 4-D attention output tensors, with the heads last layout + (`[batch, seqlen, heads, head_dim]`) + + lse_tensors (List[Tensor]): List of 3-D logsumexp tensors, with the heads last layout + (`[batch, seqlen, heads]`) + + torch_compile (bool): Attempt to use `torch.compile` to fuse the underlying elementwise + operations. Default: True. + + use_autograd_fix (bool): fix backpropagation by using a custom autograd function. Only + compatible with fused attention operations (Flash/FMHA/FNA), only as long as the inputs + of this function are (views) of outputs from said attention operation. + NATTEN's tests (tests/test_attn_merge.py) only verify correctness for when using + attention operations from NATTEN. Integration for non-NATTEN ops must be verified by the + end user. + This must be disabled when using unfused Attention, which includes Flex without + torch.compile. Default: True. + + Returns: + output (Tensor): merged attention output. + + logsumexp (Tensor): updated logsumexp. + """ + + if len(outputs) < 2: + raise ValueError("`merge_attentions` expects at least two tensors.") + + if len(outputs) != len(lse_tensors): + raise ValueError( + "`merge_attentions` expected number of outputs and LSE tensors to match, " + f"got {len(outputs)=} != {len(lse_tensors)}." + ) + assert len(outputs) == len(lse_tensors) + + requires_grad = outputs[0].requires_grad + shape = outputs[0].shape + + for i, (output, lse) in enumerate(zip(outputs, lse_tensors)): + if output.dim() != 4 or not output.is_contiguous(): + raise ValueError( + "Output tensors must be rank-4 tensors with (batch, seq, heads, dim), " + f"but got output {i} with rank={output.dim()}." + ) + + if output.shape != shape: + raise ValueError( + f"Output tensors must must match in shape, but got output {i} " + f"with shape={output.shape}." + ) + + if lse.dim() != 3: + raise ValueError( + "LSE tensors must be rank-3 tensors with (batch, seq, heads)" + f"but got LSE {i} with rank={lse.dim()}." + ) + + if lse.shape != shape[:3]: + raise ValueError( + f"LSE tensors must must match outputs in shape except last dim " + f"({shape=}), but got LSE {i} with shape={lse.shape}." + ) + + if output.requires_grad and not requires_grad: + raise ValueError( + "Either all attentions must require grad, or none of them." + ) + + # This path is the correct way to do backward pass, but since we can't have lists as inputs to + # autograd functions, we're forced to specialize it for 2-way for now. + if use_autograd_fix: + merged_output, merged_lse = MergeAttentionsAutogradFn.apply( + *outputs, *lse_tensors, torch_compile + ) + return merged_output, merged_lse + + return _merge_attentions_op( + [output.contiguous() for output in outputs], + [lse.contiguous() for lse in lse_tensors], + torch_compile=torch_compile, + ) + + +__all__ = ["merge_attentions"] diff --git a/natten/torch-ext/natten/backends/__init__.py b/natten/torch-ext/natten/backends/__init__.py new file mode 100644 index 00000000..84b8192a --- /dev/null +++ b/natten/torch-ext/natten/backends/__init__.py @@ -0,0 +1,254 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + + +from typing import List + +from ..utils import log + +logger = log.get_logger(__name__) + +import torch # noqa: F401 +from torch import Tensor + +from ..backends.blackwell_fmha import cutlass_blackwell_fmha +from ..backends.blackwell_fna import ( + cutlass_blackwell_fna_generic, + na1d_cutlass_blackwell_fna, + na2d_cutlass_blackwell_fna, + na3d_cutlass_blackwell_fna, +) +from ..backends.configs import ( + get_bwd_configs_for_cutlass_blackwell_fmha, + get_bwd_configs_for_cutlass_blackwell_fna, + get_bwd_configs_for_cutlass_fmha, + get_bwd_configs_for_cutlass_fna, + get_bwd_configs_for_cutlass_hopper_fmha, + get_bwd_configs_for_cutlass_hopper_fna, + get_configs_for_cutlass_blackwell_fmha, + get_configs_for_cutlass_blackwell_fna, + get_configs_for_cutlass_fmha, + get_configs_for_cutlass_fna, + get_configs_for_cutlass_hopper_fmha, + get_configs_for_cutlass_hopper_fna, + get_configs_for_flex_fmha, + get_configs_for_flex_fna, +) +from ..backends.configs.checks import ( + can_run_cutlass_blackwell_fmha, + can_run_cutlass_blackwell_fna, + can_run_cutlass_fna, + can_run_cutlass_hopper_fmha, + can_run_cutlass_hopper_fna, + can_run_flex_attention, +) +from ..backends.flex import ( + flex_fmha, + flex_fna_generic, + na1d_flex, + na2d_flex, + na3d_flex, +) +from ..backends.fmha import can_run_cutlass_fmha, cutlass_fmha +from ..backends.fna import ( + cutlass_fna_generic, + na1d_cutlass_fna, + na2d_cutlass_fna, + na3d_cutlass_fna, +) +from ..backends.hopper_fmha import cutlass_hopper_fmha +from ..backends.hopper_fna import ( + cutlass_hopper_fna_generic, + na1d_cutlass_hopper_fna, + na2d_cutlass_hopper_fna, + na3d_cutlass_hopper_fna, +) + + +def choose_backend( + query: Tensor, key: Tensor, value: Tensor, torch_compile: bool +) -> str: + if can_run_cutlass_blackwell_fna(query, key, value): + logger.debug("Backend not set; picked Blackwell FNA kernel.") + return "blackwell-fna" + + if can_run_cutlass_hopper_fna(query, key, value): + logger.debug("Backend not set; picked Hopper FNA kernel.") + return "hopper-fna" + + if can_run_cutlass_fna(query, key, value): + logger.debug("Backend not set; picked CUTLASS (2.X) FNA kernel.") + return "cutlass-fna" + + if can_run_flex_attention(query, key, value, torch_compile=torch_compile): + logger.debug("Backend not set; picked Flex Attention kernel.") + return "flex-fna" + + raise NotImplementedError( + "NATTEN could not find a suitable backend for this use case. " + "Run with NATTEN_LOG_LEVEL=DEBUG to find out why." + ) + + +def choose_fmha_backend( + query: Tensor, + key: Tensor, + value: Tensor, + is_causal: bool, + is_varlen: bool, + torch_compile: bool, +) -> str: + if can_run_cutlass_blackwell_fmha( + query, key, value, is_causal=is_causal, is_varlen=is_varlen + ): + logger.debug("Backend not set; picked Blackwell FMHA kernel.") + return "blackwell-fmha" + + if can_run_cutlass_hopper_fmha( + query, key, value, is_causal=is_causal, is_varlen=is_varlen + ): + logger.debug("Backend not set; picked Hopper FMHA kernel.") + return "hopper-fmha" + + if can_run_cutlass_fmha( + query, key, value, is_causal=is_causal, is_varlen=is_varlen + ): + logger.debug("Backend not set; picked CUTLASS (2.X) FMHA kernel.") + return "cutlass-fmha" + + if can_run_flex_attention( + query, + key, + value, + is_causal=is_causal, + is_varlen=is_varlen, + torch_compile=torch_compile, + ): + logger.debug("Backend not set; picked Flex Attention kernel.") + return "flex-fmha" + + raise NotImplementedError( + "NATTEN could not find a suitable backend for this FMHA use case. " + "Run with NATTEN_LOG_LEVEL=DEBUG to find out why." + ) + + +def get_compatible_backends( + query: Tensor, key: Tensor, value: Tensor, torch_compile: bool +) -> List[str]: + compatible_backends = [] + if can_run_cutlass_blackwell_fna(query, key, value): + compatible_backends.append("blackwell-fna") + + if can_run_cutlass_hopper_fna(query, key, value): + compatible_backends.append("hopper-fna") + + if can_run_cutlass_fna(query, key, value): + compatible_backends.append("cutlass-fna") + + if can_run_flex_attention(query, key, value, torch_compile=torch_compile): + compatible_backends.append("flex-fna") + + return compatible_backends + + +def get_compatible_fmha_backends( + query: Tensor, + key: Tensor, + value: Tensor, + is_causal: bool, + is_varlen: bool, + torch_compile: bool, +) -> List[str]: + compatible_backends = [] + if can_run_cutlass_blackwell_fmha( + query, key, value, is_causal=is_causal, is_varlen=is_varlen + ): + compatible_backends.append("blackwell-fmha") + + if can_run_cutlass_hopper_fmha( + query, key, value, is_causal=is_causal, is_varlen=is_varlen + ): + compatible_backends.append("hopper-fmha") + + if can_run_cutlass_fmha( + query, key, value, is_causal=is_causal, is_varlen=is_varlen + ): + compatible_backends.append("cutlass-fmha") + + if can_run_flex_attention( + query, + key, + value, + is_causal=is_causal, + is_varlen=is_varlen, + torch_compile=torch_compile, + ): + compatible_backends.append("flex-fmha") + + return compatible_backends + + +__all__ = [ + "can_run_cutlass_fmha", + "can_run_cutlass_fna", + "can_run_cutlass_blackwell_fmha", + "can_run_cutlass_blackwell_fna", + "can_run_cutlass_hopper_fmha", + "can_run_cutlass_hopper_fna", + "can_run_flex_attention", + "cutlass_fmha", + "cutlass_fna_generic", + "na1d_cutlass_fna", + "na2d_cutlass_fna", + "na3d_cutlass_fna", + "cutlass_blackwell_fmha", + "cutlass_blackwell_fna_generic", + "cutlass_hopper_fmha", + "cutlass_hopper_fna_generic", + "na1d_cutlass_blackwell_fna", + "na2d_cutlass_blackwell_fna", + "na3d_cutlass_blackwell_fna", + "flex_fmha", + "flex_fna_generic", + "na1d_flex", + "na2d_flex", + "na3d_flex", + "na1d_cutlass_hopper_fna", + "na2d_cutlass_hopper_fna", + "na3d_cutlass_hopper_fna", + "get_bwd_configs_for_cutlass_fmha", + "get_bwd_configs_for_cutlass_fna", + "get_bwd_configs_for_cutlass_blackwell_fmha", + "get_bwd_configs_for_cutlass_blackwell_fna", + "get_configs_for_cutlass_blackwell_fmha", + "get_configs_for_cutlass_blackwell_fna", + "get_configs_for_cutlass_fmha", + "get_configs_for_cutlass_fna", + "get_configs_for_cutlass_hopper_fmha", + "get_bwd_configs_for_cutlass_hopper_fmha", + "get_configs_for_cutlass_hopper_fna", + "get_bwd_configs_for_cutlass_hopper_fna", + "get_configs_for_flex_fmha", + "get_configs_for_flex_fna", +] diff --git a/natten/torch-ext/natten/backends/blackwell_fmha.py b/natten/torch-ext/natten/backends/blackwell_fmha.py new file mode 100644 index 00000000..c2816300 --- /dev/null +++ b/natten/torch-ext/natten/backends/blackwell_fmha.py @@ -0,0 +1,254 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# +import functools +from typing import Optional, Tuple, Union + +import torch +from torch import Tensor +from torch.amp import custom_bwd, custom_fwd +from torch.autograd import Function + +amp_fwd = functools.partial(custom_fwd, device_type="cuda") +amp_bwd = functools.partial(custom_bwd, device_type="cuda") + +from .._libnatten import blackwell_fmha_backward, blackwell_fmha_forward +from ..backends.configs.checks import can_run_cutlass_blackwell_fmha +from ..backends.configs.cutlass_blackwell import ( + check_cutlass_blackwell_fmha_backward_config, + check_cutlass_blackwell_fmha_forward_config, +) +from ..types import ( + CutlassBlackwellFmhaBackwardConfigType, + CutlassBlackwellFmhaForwardConfigType, + NoneType, +) +from ..utils import log +from ..utils.checks import fmha_tensor_checks, varlen_tensor_checks + +logger = log.get_logger(__name__) + + +class CutlassBlackwellFmhaAutogradFn(Function): + @staticmethod + @amp_fwd + def forward( + ctx, + query: Tensor, + key: Tensor, + value: Tensor, + is_causal: bool, + scale: float, + forward_config: CutlassBlackwellFmhaForwardConfigType, + backward_config: CutlassBlackwellFmhaBackwardConfigType, + run_persistent_kernel: bool, + cumulative_seqlen_Q: Optional[Tensor], + cumulative_seqlen_KV: Optional[Tensor], + max_seqlen_Q: int, + max_seqlen_KV: int, + ) -> Tuple[Tensor, Tensor]: + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + + q_tile_size, kv_tile_size = forward_config + + output, logsumexp = blackwell_fmha_forward( + query, + key, + value, + is_causal, + scale, + q_tile_size, + kv_tile_size, + run_persistent_kernel, + cumulative_seqlen_Q, + cumulative_seqlen_KV, + max_seqlen_Q, + max_seqlen_KV, + ) + + ctx.save_for_backward( + query, + key, + value, + logsumexp, + output, + cumulative_seqlen_Q, + cumulative_seqlen_KV, + ) + ctx.scale = scale + ctx.is_causal = is_causal + ctx.max_seqlen_Q = max_seqlen_Q + ctx.max_seqlen_KV = max_seqlen_KV + ctx.backward_config = backward_config + # Always record determinism behavior during forward pass (forward pass itself is + # deterministic anyway). + # Determinism could be limited to part of the program, which means during forward pass + # it'll be true, but on .backward() call, if it's been turned off, it will stay off when we + # get to this operation's backward call. + ctx.deterministic = torch.are_deterministic_algorithms_enabled() + + return output, logsumexp + + @staticmethod + @amp_bwd + def backward(ctx, grad_out: Tensor, grad_lse: Tensor) -> Tuple[ + Tensor, + Tensor, + Tensor, + NoneType, + NoneType, + NoneType, + NoneType, + NoneType, + # varlen + NoneType, + NoneType, + NoneType, + NoneType, + ]: + ( + query, + key, + value, + logsumexp, + output, + cumulative_seqlen_Q, + cumulative_seqlen_KV, + ) = ctx.saved_tensors + d_output = grad_out.contiguous() # noqa: F841 + + q_tile_size, k_tile_size = ctx.backward_config + + d_query, d_key, d_value = blackwell_fmha_backward( + query, + key, + value, + output, + d_output, + logsumexp, + ctx.is_causal, + ctx.scale, + q_tile_size, + k_tile_size, + cumulative_seqlen_Q, + cumulative_seqlen_KV, + ctx.max_seqlen_Q, + ctx.max_seqlen_KV, + ctx.deterministic, + ) + + return ( + d_query, + d_key, + d_value, + None, + None, + None, + None, + None, + None, + None, + None, + None, + ) + + +def cutlass_blackwell_fmha( + query: Tensor, + key: Tensor, + value: Tensor, + is_causal: bool = False, + scale: Optional[float] = None, + q_tile_size: Optional[int] = None, + kv_tile_size: Optional[int] = None, + run_persistent_kernel: bool = False, + backward_q_tile_size: Optional[int] = None, + backward_kv_tile_size: Optional[int] = None, + return_lse: bool = False, + # varlen parameters + cumulative_seqlen_Q: Optional[Tensor] = None, + cumulative_seqlen_KV: Optional[Tensor] = None, + max_seqlen_Q: int = 0, + max_seqlen_KV: int = 0, +) -> Union[Tensor, Tuple[Tensor, Tensor]]: + + fmha_tensor_checks( + query, + key, + value, + must_match_head_dims=True, + supports_gqa_mqa=True, + backend_name="Blackwell FMHA", + ) + + ( + cumulative_seqlen_Q, + cumulative_seqlen_KV, + max_seqlen_Q, + max_seqlen_KV, + ) = varlen_tensor_checks( + query=query, + key=key, + value=value, + cumulative_seqlen_Q=cumulative_seqlen_Q, + cumulative_seqlen_KV=cumulative_seqlen_KV, + max_seqlen_Q=max_seqlen_Q, + max_seqlen_KV=max_seqlen_KV, + ) + is_varlen = cumulative_seqlen_Q is not None + + assert can_run_cutlass_blackwell_fmha( + query, key, value, is_causal=is_causal, is_varlen=is_varlen, raise_error=True + ) + + forward_config = check_cutlass_blackwell_fmha_forward_config( + input_tensor=query, q_tile_size=q_tile_size, kv_tile_size=kv_tile_size + ) + backward_config = check_cutlass_blackwell_fmha_backward_config( + input_tensor=query, + q_tile_size=backward_q_tile_size, + kv_tile_size=backward_kv_tile_size, + ) + + scale = scale or query.shape[-1] ** -0.5 + + output, lse = CutlassBlackwellFmhaAutogradFn.apply( + query, + key, + value, + is_causal, + scale, + forward_config, + backward_config, + run_persistent_kernel, + cumulative_seqlen_Q, + cumulative_seqlen_KV, + max_seqlen_Q, + max_seqlen_KV, + ) + + if return_lse: + return output, lse + + return output diff --git a/natten/torch-ext/natten/backends/blackwell_fna.py b/natten/torch-ext/natten/backends/blackwell_fna.py new file mode 100644 index 00000000..184faa0b --- /dev/null +++ b/natten/torch-ext/natten/backends/blackwell_fna.py @@ -0,0 +1,500 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# +import functools +from typing import Optional, Tuple, Union + +import torch +from torch import Tensor +from torch.amp import custom_bwd, custom_fwd +from torch.autograd import Function + +amp_fwd = functools.partial(custom_fwd, device_type="cuda") +amp_bwd = functools.partial(custom_bwd, device_type="cuda") + +from .._libnatten import ( + blackwell_na1d_backward, + blackwell_na1d_forward, + blackwell_na2d_backward, + blackwell_na2d_forward, + blackwell_na3d_backward, + blackwell_na3d_forward, +) +from ..backends.configs.checks import can_run_cutlass_blackwell_fna +from ..backends.configs.cutlass_blackwell import ( + check_cutlass_blackwell_fna_backward_config, + check_cutlass_blackwell_fna_forward_config, +) +from ..token_permute import token_permute_operation, token_unpermute_operation +from ..types import ( + CausalArg1DTypeOrDed, + CausalArg2DTypeOrDed, + CausalArg3DTypeOrDed, + CausalArgType, + CausalArgTypeOrDed, + CutlassBlackwellFnaBackwardConfigType, + CutlassBlackwellFnaForwardConfigType, + Dimension1DType, + Dimension1DTypeOrDed, + Dimension2DType, + Dimension2DTypeOrDed, + Dimension3DType, + Dimension3DTypeOrDed, + DimensionType, + DimensionTypeOrDed, + NoneType, +) +from ..utils.checks import ( + check_all_args, + check_args_against_input, + na_tensor_checks, +) + + +def make_cutlass_blackwell_fna_autograd_fn(na_dim): + assert na_dim in [1, 2, 3] + + FORWARD_OPS = { + 1: blackwell_na1d_forward, + 2: blackwell_na2d_forward, + 3: blackwell_na3d_forward, + } + + BACKWARD_OPS = { + 1: blackwell_na1d_backward, + 2: blackwell_na2d_backward, + 3: blackwell_na3d_backward, + } + + class CutlassBlackwellFnaGenericAutogradFn(Function): + @staticmethod + @amp_fwd + def forward( + ctx, + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size: DimensionType, + stride: DimensionType, + dilation: DimensionType, + is_causal: CausalArgType, + scale: float, + forward_config: CutlassBlackwellFnaForwardConfigType, + backward_config: CutlassBlackwellFnaBackwardConfigType, + run_persistent_kernel: bool, + ) -> Tuple[Tensor, Tensor]: + kernel_size, stride, dilation, is_causal = check_all_args( + na_dim, kernel_size, stride, dilation, is_causal + ) + + q_tile_shape, kv_tile_shape = forward_config + + # Token permute begin + query_perm, qkv_shape, q_shape = token_permute_operation( + query, + tile_shape=q_tile_shape, + dilation=dilation, + flip_tiled_dims=True, + ) + key_perm, _, k_shape = token_permute_operation( + key, tile_shape=kv_tile_shape, dilation=dilation, flip_tiled_dims=True + ) + value_perm, _, v_shape = token_permute_operation( + value, tile_shape=kv_tile_shape, dilation=dilation, flip_tiled_dims=True + ) + + assert k_shape == v_shape + kv_shape = k_shape + # Token permute end + + query_perm = query_perm.contiguous() + key_perm = key_perm.contiguous() + value_perm = value_perm.contiguous() + + output_perm, logsumexp_perm = FORWARD_OPS[na_dim]( + query_perm, + key_perm, + value_perm, + kernel_size, + stride, + dilation, + is_causal, + scale, + q_shape, + kv_shape, + qkv_shape, + q_tile_shape, + kv_tile_shape, + run_persistent_kernel, + ) + + # Token un-permute begin + output = token_unpermute_operation( + output_perm, + token_layout_shape=qkv_shape, + tile_shape=q_tile_shape, + dilation=dilation, + flip_tiled_dims=True, + ) + logsumexp = token_unpermute_operation( + logsumexp_perm.unsqueeze(-1), + token_layout_shape=qkv_shape, + tile_shape=q_tile_shape, + dilation=dilation, + flip_tiled_dims=True, + ).squeeze(-1) + # Token un-permute end + + ctx.save_for_backward(query, key, value, logsumexp, output) + ctx.kernel_size = kernel_size + ctx.stride = stride + ctx.dilation = dilation + ctx.is_causal = is_causal + ctx.scale = scale + ctx.backward_config = backward_config + # Always record determinism behavior during forward pass (forward pass itself is + # deterministic anyway). + # Determinism could be limited to part of the program, which means during forward pass + # it'll be true, but on .backward() call, if it's been turned off, it will stay off when we + # get to this operation's backward call. + ctx.deterministic = torch.are_deterministic_algorithms_enabled() + + return output, logsumexp + + @staticmethod + @amp_bwd + def backward(ctx, d_output: Tensor, d_lse: Tensor) -> Tuple[ + Tensor, + Tensor, + Tensor, + NoneType, + NoneType, + NoneType, + NoneType, + NoneType, + NoneType, + NoneType, + NoneType, + ]: + query, key, value, logsumexp, output = ctx.saved_tensors + kernel_size, stride, dilation, is_causal, scale = ( + ctx.kernel_size, + ctx.stride, + ctx.dilation, + ctx.is_causal, + ctx.scale, + ) + + q_tile_shape, kv_tile_shape = ctx.backward_config + + if ctx.deterministic: + raise RuntimeError( + "Blackwell FNA backward pass does not have a deterministic mode, " + "but PyTorch's deterministic algorithms were enabled. To proceed, " + "you must either disable torch's deterministic mode, or choose a " + "different backend." + ) + + # Token permute begin + + query_perm, qkv_shape, q_shape = token_permute_operation( + query, tile_shape=q_tile_shape, dilation=dilation, flip_tiled_dims=True + ) + output_perm, _, o_shape = token_permute_operation( + output, tile_shape=q_tile_shape, dilation=dilation, flip_tiled_dims=True + ) + d_output_perm, _, d_o_shape = token_permute_operation( + d_output, + tile_shape=q_tile_shape, + dilation=dilation, + flip_tiled_dims=True, + ) + logsumexp_perm, _, _ = token_permute_operation( + logsumexp.unsqueeze(-1), + tile_shape=q_tile_shape, + dilation=dilation, + flip_tiled_dims=True, + ) + key_perm, _, k_shape = token_permute_operation( + key, tile_shape=kv_tile_shape, dilation=dilation, flip_tiled_dims=True + ) + value_perm, _, v_shape = token_permute_operation( + value, tile_shape=kv_tile_shape, dilation=dilation, flip_tiled_dims=True + ) + + assert q_shape == o_shape == d_o_shape + assert k_shape == v_shape + kv_shape = k_shape + # Token permute end + + query_perm = query_perm.contiguous() + key_perm = key_perm.contiguous() + value_perm = value_perm.contiguous() + output_perm = output_perm.contiguous() + d_output_perm = d_output_perm.contiguous() + logsumexp_perm = logsumexp_perm.squeeze(-1) + + d_query_perm, d_key_perm, d_value_perm = BACKWARD_OPS[na_dim]( + query_perm, + key_perm, + value_perm, + output_perm, + d_output_perm, + logsumexp_perm, + kernel_size, + stride, + dilation, + is_causal, + scale, + q_shape, + kv_shape, + qkv_shape, + q_tile_shape, + kv_tile_shape, + ) + + # Token un-permute begin + d_query = token_unpermute_operation( + d_query_perm, + token_layout_shape=qkv_shape, + tile_shape=q_tile_shape, + dilation=dilation, + flip_tiled_dims=True, + ) + d_key = token_unpermute_operation( + d_key_perm, + token_layout_shape=qkv_shape, + tile_shape=kv_tile_shape, + dilation=dilation, + flip_tiled_dims=True, + ) + d_value = token_unpermute_operation( + d_value_perm, + token_layout_shape=qkv_shape, + tile_shape=kv_tile_shape, + dilation=dilation, + flip_tiled_dims=True, + ) + # Token un-permute end + + assert d_query.shape == query.shape + assert d_key.shape == key.shape + assert d_value.shape == value.shape + + return ( + d_query, + d_key, + d_value, + None, + None, + None, + None, + None, + None, + None, + None, + ) + + return CutlassBlackwellFnaGenericAutogradFn + + +CutlassBlackwellFna1DAutogradFn = make_cutlass_blackwell_fna_autograd_fn(1) +CutlassBlackwellFna2DAutogradFn = make_cutlass_blackwell_fna_autograd_fn(2) +CutlassBlackwellFna3DAutogradFn = make_cutlass_blackwell_fna_autograd_fn(3) + + +CutlassBlackwellFNAAutogradFns = { + 1: CutlassBlackwellFna1DAutogradFn, + 2: CutlassBlackwellFna2DAutogradFn, + 3: CutlassBlackwellFna3DAutogradFn, +} + + +def cutlass_blackwell_fna_generic( + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size: DimensionTypeOrDed, + stride: DimensionTypeOrDed = 1, + dilation: DimensionTypeOrDed = 1, + is_causal: Optional[CausalArgTypeOrDed] = False, + scale: Optional[float] = None, + q_tile_shape: Optional[DimensionType] = None, + kv_tile_shape: Optional[DimensionType] = None, + backward_q_tile_shape: Optional[DimensionType] = None, + backward_kv_tile_shape: Optional[DimensionType] = None, + run_persistent_kernel: bool = False, + return_lse: bool = False, +) -> Union[Tensor, Tuple[Tensor, Tensor]]: + + na_tensor_checks( + query, key, value, must_match_head_dims=True, supports_gqa_mqa=True + ) + + assert can_run_cutlass_blackwell_fna(query, key, value, raise_error=True) + + na_dim = query.dim() - 3 # batch, heads, head_dim + + kernel_size, stride, dilation, is_causal = check_all_args( + na_dim, kernel_size, stride, dilation, is_causal + ) + + check_args_against_input( + query, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + ) + + forward_config = check_cutlass_blackwell_fna_forward_config( + input_tensor=query, q_tile_shape=q_tile_shape, kv_tile_shape=kv_tile_shape + ) + + backward_config = check_cutlass_blackwell_fna_backward_config( + input_tensor=query, + q_tile_shape=backward_q_tile_shape, + kv_tile_shape=backward_kv_tile_shape, + ) + + scale = scale or query.shape[-1] ** -0.5 + + output, lse = CutlassBlackwellFNAAutogradFns[na_dim].apply( + query, + key, + value, + kernel_size, + stride, + dilation, + is_causal, + scale, + forward_config, + backward_config, + run_persistent_kernel, + ) + + if return_lse: + return output, lse + + return output + + +def na1d_cutlass_blackwell_fna( + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size: Dimension1DTypeOrDed, + stride: Dimension1DTypeOrDed = 1, + dilation: Dimension1DTypeOrDed = 1, + is_causal: Optional[CausalArg1DTypeOrDed] = False, + scale: Optional[float] = None, + q_tile_shape: Optional[Dimension1DType] = None, + kv_tile_shape: Optional[Dimension1DType] = None, + backward_q_tile_shape: Optional[Dimension1DType] = None, + backward_kv_tile_shape: Optional[Dimension1DType] = None, + run_persistent_kernel: bool = False, + return_lse: bool = False, +) -> Union[Tensor, Tuple[Tensor, Tensor]]: + return cutlass_blackwell_fna_generic( + query=query, + key=key, + value=value, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + scale=scale, + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + backward_q_tile_shape=backward_q_tile_shape, + backward_kv_tile_shape=backward_kv_tile_shape, + run_persistent_kernel=run_persistent_kernel, + return_lse=return_lse, + ) + + +def na2d_cutlass_blackwell_fna( + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size: Dimension2DTypeOrDed, + stride: Dimension2DTypeOrDed = 1, + dilation: Dimension2DTypeOrDed = 1, + is_causal: Optional[CausalArg2DTypeOrDed] = False, + scale: Optional[float] = None, + q_tile_shape: Optional[Dimension2DType] = None, + kv_tile_shape: Optional[Dimension2DType] = None, + backward_q_tile_shape: Optional[Dimension2DType] = None, + backward_kv_tile_shape: Optional[Dimension2DType] = None, + run_persistent_kernel: bool = False, + return_lse: bool = False, +) -> Union[Tensor, Tuple[Tensor, Tensor]]: + return cutlass_blackwell_fna_generic( + query=query, + key=key, + value=value, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + scale=scale, + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + backward_q_tile_shape=backward_q_tile_shape, + backward_kv_tile_shape=backward_kv_tile_shape, + run_persistent_kernel=run_persistent_kernel, + return_lse=return_lse, + ) + + +def na3d_cutlass_blackwell_fna( + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size: Dimension3DTypeOrDed, + stride: Dimension3DTypeOrDed = 1, + dilation: Dimension3DTypeOrDed = 1, + is_causal: Optional[CausalArg3DTypeOrDed] = False, + scale: Optional[float] = None, + q_tile_shape: Optional[Dimension3DType] = None, + kv_tile_shape: Optional[Dimension3DType] = None, + backward_q_tile_shape: Optional[Dimension3DType] = None, + backward_kv_tile_shape: Optional[Dimension3DType] = None, + run_persistent_kernel: bool = False, + return_lse: bool = False, +) -> Union[Tensor, Tuple[Tensor, Tensor]]: + return cutlass_blackwell_fna_generic( + query=query, + key=key, + value=value, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + scale=scale, + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + backward_q_tile_shape=backward_q_tile_shape, + backward_kv_tile_shape=backward_kv_tile_shape, + run_persistent_kernel=run_persistent_kernel, + return_lse=return_lse, + ) diff --git a/natten/torch-ext/natten/backends/configs/__init__.py b/natten/torch-ext/natten/backends/configs/__init__.py new file mode 100644 index 00000000..732097f2 --- /dev/null +++ b/natten/torch-ext/natten/backends/configs/__init__.py @@ -0,0 +1,584 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + + +from typing import List + +from ...utils import log + +logger = log.get_logger(__name__) + +import torch # noqa: F401 +from torch import Tensor + +from ...backends.configs.checks import ( + can_run_cutlass_blackwell_fmha, + can_run_cutlass_blackwell_fna, + can_run_cutlass_fmha, + can_run_cutlass_fna, + can_run_cutlass_hopper_fmha, + can_run_cutlass_hopper_fna, + can_run_flex_attention, +) +from ...backends.configs.cutlass import ( + get_all_tile_shapes_backward as get_all_cutlass_fna_backward_configs, + get_all_tile_shapes_forward as get_all_cutlass_fna_forward_configs, + get_all_tile_sizes_backward as get_all_cutlass_fmha_backward_configs, + get_all_tile_sizes_forward as get_all_cutlass_fmha_forward_configs, +) +from ...backends.configs.cutlass_blackwell import ( + get_all_backward_configs as get_all_blackwell_fna_backward_configs, + get_all_fmha_backward_configs as get_all_blackwell_fmha_backward_configs, + get_all_fmha_forward_configs as get_all_blackwell_fmha_forward_configs, + get_all_forward_configs as get_all_blackwell_fna_forward_configs, +) +from ...backends.configs.cutlass_hopper import ( + get_all_backward_configs as get_all_hopper_fna_backward_configs, + get_all_fmha_backward_configs as get_all_hopper_fmha_backward_configs, + get_all_fmha_forward_configs as get_all_hopper_fmha_forward_configs, + get_all_forward_configs as get_all_hopper_fna_forward_configs, +) +from ...backends.configs.flex import ( + get_all_tile_shapes_forward as get_all_flex_fna_forward_configs, + get_all_tile_sizes_forward as get_all_flex_fmha_forward_configs, +) +from ...types import ( + CutlassBlackwellFmhaBackwardConfigType, + CutlassBlackwellFmhaForwardConfigType, + CutlassBlackwellFnaBackwardConfigType, + CutlassBlackwellFnaForwardConfigType, + CutlassFmhaBackwardConfigType, + CutlassFmhaForwardConfigType, + CutlassFnaBackwardConfigType, + CutlassFnaForwardConfigType, + CutlassHopperFmhaBackwardConfigType, + CutlassHopperFmhaForwardConfigType, + CutlassHopperFnaBackwardConfigType, + CutlassHopperFnaForwardConfigType, + FlexFmhaForwardConfigType, + FlexFnaForwardConfigType, +) + +### CUTLASS Blackwell kernels + + +def get_configs_for_cutlass_blackwell_fmha( + query: Tensor, + key: Tensor, + value: Tensor, +) -> List[CutlassBlackwellFmhaForwardConfigType]: + """Returns Blackwell FMHA configurations compatible with input tensors, if any. + + Checks first if a CUDA tensor, and on a Blackwell datacenter GPU (SM100; compute capability + 10.0), and if so, returns *forward pass* configurations compatible with the tensor dtype and + head dim. + + Each configuration for this operation is a tuple of two integers: `(q_tile_size, + kv_tile_size)`. These are arguments to [natten.attention][natten.attention]. + + Args: + query: Query tensor matching the shape, dtype, and device of your use case. + key: Key tensor matching the shape, dtype, and device of your use case. + value: Value tensor matching the shape, dtype, and device of your use case. + + Returns: + (List[Tuple[int, int]]): List of tuples of two integers corresponding to query and KV tile + sizes. + """ + + if not can_run_cutlass_blackwell_fmha( + query=query, + key=key, + value=value, + is_causal=False, + is_varlen=False, + raise_error=False, + ): + return [] + + return get_all_blackwell_fmha_forward_configs(input_tensor=query) + + +def get_bwd_configs_for_cutlass_blackwell_fmha( + query: Tensor, + key: Tensor, + value: Tensor, +) -> List[CutlassBlackwellFmhaBackwardConfigType]: + """Returns Blackwell FMHA backward pass configurations compatible with input tensors, if any. + + Checks first if a CUDA tensor, and on a Blackwell datacenter GPU (SM100; compute capability + 10.0), and if so, returns *backward pass* configurations compatible with the tensor dtype and + head dim. + + Each configuration for this operation is a tuple of two integers: `(q_tile_size, + kv_tile_size)`. These are arguments to [natten.attention][natten.attention]. + + Args: + query: Query tensor matching the shape, dtype, and device of your use case. + key: Key tensor matching the shape, dtype, and device of your use case. + value: Value tensor matching the shape, dtype, and device of your use case. + + Returns: + (List[Tuple[int, int]]): List of tuples of two integers corresponding to query and KV tile + sizes. + """ + + if not can_run_cutlass_blackwell_fmha( + query=query, + key=key, + value=value, + is_causal=False, + is_varlen=False, + raise_error=False, + ): + return [] + + return get_all_blackwell_fmha_backward_configs(input_tensor=query) + + +def get_configs_for_cutlass_blackwell_fna( + query: Tensor, + key: Tensor, + value: Tensor, +) -> List[CutlassBlackwellFnaForwardConfigType]: + """Returns Blackwell FNA configurations compatible with input tensors, if any. + + Checks first if a CUDA tensor, and on a Blackwell datacenter GPU (SM100; compute capability + 10.0), and if so, returns *forward pass* configurations compatible with the tensor dtype and + head dim, and according to the rank of the token layout (1D/2D/3D). + + Each configuration for this operation is a tuple of two integer tuples: `(q_tile_shape, + kv_tile_shape)`. These are arguments to [natten.na1d][natten.na1d], [natten.na2d][natten.na2d], + and [natten.na3d][natten.na3d]. + + Args: + query: Query tensor matching the shape, dtype, and device of your use case. + key: Key tensor matching the shape, dtype, and device of your use case. + value: Value tensor matching the shape, dtype, and device of your use case. + + Returns: + (List[Tuple[tuple, tuple]]): List of tuples of two integer tuples corresponding to query + and KV tile *shapes*. + """ + if not can_run_cutlass_blackwell_fna( + query=query, key=key, value=value, raise_error=False + ): + return [] + + return get_all_blackwell_fna_forward_configs(input_tensor=query) + + +def get_bwd_configs_for_cutlass_blackwell_fna( + query: Tensor, + key: Tensor, + value: Tensor, +) -> List[CutlassBlackwellFnaBackwardConfigType]: + """Returns Blackwell FNA backward pass configurations compatible with input tensors, if any. + + Checks first if a CUDA tensor, and on a Blackwell datacenter GPU (SM100; compute capability + 10.0), and if so, returns *backward pass* configurations compatible with the tensor dtype and + head dim, and according to the rank of the token layout (1D/2D/3D). + + Each configuration for this operation is a tuple of two integer tuples: `(q_tile_shape, + kv_tile_shape)`. These are arguments to [natten.na1d][natten.na1d], [natten.na2d][natten.na2d], + and [natten.na3d][natten.na3d]. + + Args: + query: Query tensor matching the shape, dtype, and device of your use case. + key: Key tensor matching the shape, dtype, and device of your use case. + value: Value tensor matching the shape, dtype, and device of your use case. + + Returns: + (List[Tuple[tuple, tuple]]): List of tuples of two integer tuples corresponding to query + and KV tile *shapes*. + """ + if not can_run_cutlass_blackwell_fna( + query=query, key=key, value=value, raise_error=False + ): + return [] + + return get_all_blackwell_fna_backward_configs(input_tensor=query) + + +### CUTLASS Hopper kernels + + +def get_configs_for_cutlass_hopper_fmha( + query: Tensor, + key: Tensor, + value: Tensor, +) -> List[CutlassHopperFmhaForwardConfigType]: + """Returns Hopper FMHA configurations compatible with input tensors, if any. + + Checks first if a CUDA tensor, and on a Hopper GPU (SM90; compute capability 9.0), and if so, + returns *forward pass* configurations compatible with the tensor dtype and head dim. + + Each configuration for this operation is a tuple of one integer tuple, and another integer: + `((q_tile_size, kv_tile_size), kernel_schedule)`. These are arguments to + [natten.attention][natten.attention]. + `kernel_schedule` is specific to Hopper FNA/FMHA only. + + Args: + query: Query tensor matching the shape, dtype, and device of your use case. + key: Key tensor matching the shape, dtype, and device of your use case. + value: Value tensor matching the shape, dtype, and device of your use case. + + Returns: + (List[Tuple[Tuple[int, int], KernelSchedule]]): List of tuples of one tuple of two integers + corresponding to query and KV tile sizes, and a kernel schedule enum type. + """ + if not can_run_cutlass_hopper_fmha( + query=query, + key=key, + value=value, + is_causal=False, + is_varlen=False, + raise_error=False, + ): + return [] + + return get_all_hopper_fmha_forward_configs(input_tensor=query) + + +def get_bwd_configs_for_cutlass_hopper_fmha( + query: Tensor, + key: Tensor, + value: Tensor, +) -> List[CutlassHopperFmhaBackwardConfigType]: + """Returns Hopper FMHA backward pass configurations compatible with input tensors, if any. + + Checks first if a CUDA tensor, and on a Hopper GPU (SM90; compute capability 9.0), and if so, + returns *backward pass* configurations compatible with the tensor dtype and head dim. + + Each configuration for this operation is an integer tuple: + `(backward_q_tile_size, backward_kv_tile_size)`. These are arguments to + [natten.attention][natten.attention]. + + Note that unlike forward pass, kernel schedule is not part of the configuration. All backward + pass kernels are persistent warp-specialized. + See CUTLASS's [example 88](https://github.com/NVIDIA/cutlass/tree/main/examples/88_hopper_fmha) + for more. + + Args: + query: Query tensor matching the shape, dtype, and device of your use case. + key: Key tensor matching the shape, dtype, and device of your use case. + value: Value tensor matching the shape, dtype, and device of your use case. + + Returns: + (List[Tuple[int, int]]): List of integer tuples corresponding to query and KV tile sizes. + """ + if not can_run_cutlass_hopper_fmha( + query=query, + key=key, + value=value, + is_causal=False, + is_varlen=False, + raise_error=False, + ): + return [] + + return get_all_hopper_fmha_backward_configs(input_tensor=query) + + +def get_configs_for_cutlass_hopper_fna( + query: Tensor, + key: Tensor, + value: Tensor, +) -> List[CutlassHopperFnaForwardConfigType]: + """Returns Hopper FNA configurations compatible with input tensors, if any. + + Checks first if a CUDA tensor, and on a Hopper GPU (SM90; compute capability 9.0), and if so, + returns *forward pass* configurations compatible with the tensor dtype and head dim. + + Each configuration for this operation is a tuple of one tuple, and another integer: + `((q_tile_shape, kv_tile_shape), kernel_schedule)`. These are arguments to + [natten.na1d][natten.na1d], [natten.na2d][natten.na2d], and [natten.na3d][natten.na3d]. + + `kernel_schedule` is specific to Hopper FNA/FMHA only. + + Args: + query: Query tensor matching the shape, dtype, and device of your use case. + key: Key tensor matching the shape, dtype, and device of your use case. + value: Value tensor matching the shape, dtype, and device of your use case. + + Returns: + (List[Tuple[Tuple[tuple, tuple], KernelSchedule]]): List of tuples of one tuple of two + shape tuples, corresponding to query and KV tile *shapes*, and a kernel schedule enum + type. + """ + if not can_run_cutlass_hopper_fna( + query=query, key=key, value=value, raise_error=False + ): + return [] + + return get_all_hopper_fna_forward_configs(input_tensor=query) + + +def get_bwd_configs_for_cutlass_hopper_fna( + query: Tensor, + key: Tensor, + value: Tensor, +) -> List[CutlassHopperFnaBackwardConfigType]: + """Returns Hopper FNA backward pass configurations compatible with input tensors, if any. + + Checks first if a CUDA tensor, and on a Hopper GPU (SM90; compute capability 9.0), and if so, + returns *backward pass* configurations compatible with the tensor dtype and head dim. + + Each configuration for this operation is a tuple of two tuples: + `(q_tile_shape, kv_tile_shape)`. These are arguments to [natten.na1d][natten.na1d], + [natten.na2d][natten.na2d], and [natten.na3d][natten.na3d]. + + Note that unlike forward pass, kernel schedule is not part of the configuration. All backward + pass kernels are persistent warp-specialized. + See CUTLASS's [example 88](https://github.com/NVIDIA/cutlass/tree/main/examples/88_hopper_fmha) + for more. + + Args: + query: Query tensor matching the shape, dtype, and device of your use case. + key: Key tensor matching the shape, dtype, and device of your use case. + value: Value tensor matching the shape, dtype, and device of your use case. + + Returns: + (List[Tuple[tuple, tuple]]): List of tuples of two shape tuples, corresponding to query and + KV tile *shapes*. + """ + if not can_run_cutlass_hopper_fna( + query=query, key=key, value=value, raise_error=False + ): + return [] + + return get_all_hopper_fna_backward_configs(input_tensor=query) + + +### CUTLASS 2.X kernels + + +def get_configs_for_cutlass_fmha( + query: Tensor, + key: Tensor, + value: Tensor, +) -> List[CutlassFmhaForwardConfigType]: + """Returns CUTLASS FMHA configurations compatible with input tensors, if any. + + Checks first if a CUDA tensor, and on a device with compute capability >= 5.0, and if so, + returns *forward pass* configurations compatible with the specific compute capability, tensor + dtype and head dim. + + Each configuration for this operation is a tuple of two integers: `(q_tile_size, + kv_tile_size)`. These are arguments to [natten.attention][natten.attention]. + + Args: + query: Query tensor matching the shape, dtype, and device of your use case. + key: Key tensor matching the shape, dtype, and device of your use case. + value: Value tensor matching the shape, dtype, and device of your use case. + + Returns: + (List[Tuple[int, int]]): List of tuples of two integers corresponding to query and KV tile + sizes. + """ + if not can_run_cutlass_fmha( + query=query, + key=key, + value=value, + is_causal=False, + is_varlen=False, + raise_error=False, + ): + return [] + + return get_all_cutlass_fmha_forward_configs(input_tensor=query) + + +def get_bwd_configs_for_cutlass_fmha( + query: Tensor, + key: Tensor, + value: Tensor, +) -> List[CutlassFmhaBackwardConfigType]: + """Returns CUTLASS FMHA backward pass configurations compatible with input tensors, if any. + + Checks first if a CUDA tensor, and on a device with compute capability >= 5.0, and if so, + returns *backward pass* configurations compatible with the specific compute capability, tensor + dtype and head dim. + + Each configuration for this operation is a tuple of two integers: `(backward_q_tile_size, + backward_kv_tile_size)`. These are arguments to [natten.attention][natten.attention]. + + Args: + query: Query tensor matching the shape, dtype, and device of your use case. + key: Key tensor matching the shape, dtype, and device of your use case. + value: Value tensor matching the shape, dtype, and device of your use case. + + Returns: + (List[Tuple[int, int]]): List of tuples of two integers corresponding to query and KV tile + sizes in the *backward pass*. + """ + if not can_run_cutlass_fmha( + query=query, + key=key, + value=value, + is_causal=False, + is_varlen=False, + raise_error=False, + ): + return [] + + return get_all_cutlass_fmha_backward_configs( + input_tensor=key if key.shape[-1] >= value.shape[-1] else value + ) + + +def get_configs_for_cutlass_fna( + query: Tensor, + key: Tensor, + value: Tensor, +) -> List[CutlassFnaForwardConfigType]: + """Returns CUTLASS FNA configurations compatible with input tensors, if any. + + Checks first if a CUDA tensor, and on a device with compute capability >= 5.0, and if so, + returns *forward pass* configurations compatible with the specific compute capability, tensor + dtype and head dim, and according to the rank of the token layout (1D/2D/3D). + + Each configuration for this operation is a tuple of two integer tuples: `(q_tile_shape, + kv_tile_shape)`. These are arguments to [natten.na1d][natten.na1d], [natten.na2d][natten.na2d], + and [natten.na3d][natten.na3d]. + + Args: + query: Query tensor matching the shape, dtype, and device of your use case. + key: Key tensor matching the shape, dtype, and device of your use case. + value: Value tensor matching the shape, dtype, and device of your use case. + + Returns: + (List[Tuple[tuple, tuple]]): List of tuples of two integer tuples corresponding to query + and KV tile *shapes*. + """ + if not can_run_cutlass_fna(query=query, key=key, value=value, raise_error=False): + return [] + + return get_all_cutlass_fna_forward_configs(input_tensor=query) + + +def get_bwd_configs_for_cutlass_fna( + query: Tensor, + key: Tensor, + value: Tensor, +) -> List[CutlassFnaBackwardConfigType]: + """Returns CUTLASS FNA backward pass configurations compatible with input tensors, if any. + + Checks first if a CUDA tensor, and on a device with compute capability >= 5.0, and if so, + returns *backward pass* configurations compatible with the specific compute capability, tensor + dtype and head dim. + + Each configuration for this operation is a tuple of two integers: `(backward_q_tile_shape, + backward_kv_tile_shape)`. These are arguments to [natten.na1d][natten.na1d], + [natten.na2d][natten.na2d], and [natten.na3d][natten.na3d]. + + Args: + query: Query tensor matching the shape, dtype, and device of your use case. + key: Key tensor matching the shape, dtype, and device of your use case. + value: Value tensor matching the shape, dtype, and device of your use case. + + Returns: + (List[Tuple[tuple, tuple]]): List of tuples of two integer tuples corresponding to query + and KV tile *shapes* in the *backward pass*. + """ + if not can_run_cutlass_fna(query=query, key=key, value=value, raise_error=False): + return [] + + return get_all_cutlass_fna_backward_configs( + input_tensor=key if key.shape[-1] >= value.shape[-1] else value + ) + + +### Flex + + +def get_configs_for_flex_fmha( + query: Tensor, + key: Tensor, + value: Tensor, + torch_compile: bool = False, +) -> List[FlexFmhaForwardConfigType]: + """Returns Flex FMHA configurations compatible with input tensors, if any. + + Each configuration for this operation is a tuple of two integers: `(q_tile_size, + kv_tile_size)`. These are arguments to [natten.attention][natten.attention]. + Not specifying these arguments while backend is Flex will default to `q_tile_size = 64` and + `kv_tile_size = 64`. + + Args: + query: Query tensor matching the shape, dtype, and device of your use case. + key: Key tensor matching the shape, dtype, and device of your use case. + value: Value tensor matching the shape, dtype, and device of your use case. + torch_compile: Whether or not you intend to use compiled block mask and flex attention kernel. + + Returns: + (List[Tuple[int, int]]): List of tuples of two integers corresponding to query and KV tile + sizes. + """ + if not can_run_flex_attention( + query=query, + key=key, + value=value, + torch_compile=torch_compile, + raise_error=False, + ): + return [] + + return get_all_flex_fmha_forward_configs(input_tensor=query) + + +def get_configs_for_flex_fna( + query: Tensor, + key: Tensor, + value: Tensor, + torch_compile: bool = False, +) -> List[FlexFnaForwardConfigType]: + """Returns Flex FNA configurations compatible with input tensors, if any. + + Each configuration for this operation is a tuple of two integer tuples: `(q_tile_shape, + kv_tile_shape)`. These are arguments to [natten.na1d][natten.na1d], [natten.na2d][natten.na2d], + and [natten.na3d][natten.na3d]. + Not specifying these arguments while backend is Flex will default to single-dimensional tiling, + and will not use our Token Permutation approach. By explicitly specifying tile shapes, you will + automatically use our Token Permutation approach, which saves you the most compute. + + Args: + query: Query tensor matching the shape, dtype, and device of your use case. + key: Key tensor matching the shape, dtype, and device of your use case. + value: Value tensor matching the shape, dtype, and device of your use case. + torch_compile: Whether or not you intend to use compiled block mask and flex attention kernel. + + Returns: + (List[Tuple[tuple, tuple]]): List of tuples of two integer tuples corresponding to query + and KV tile *shapes*. + """ + if not can_run_flex_attention( + query=query, + key=key, + value=value, + torch_compile=torch_compile, + raise_error=False, + ): + return [] + + return get_all_flex_fna_forward_configs(input_tensor=query) diff --git a/natten/torch-ext/natten/backends/configs/checks.py b/natten/torch-ext/natten/backends/configs/checks.py new file mode 100644 index 00000000..4237d8eb --- /dev/null +++ b/natten/torch-ext/natten/backends/configs/checks.py @@ -0,0 +1,750 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# +import functools +import math + +import torch +from torch import Tensor + +from ..._environment import _IS_TORCH_COMPILE_SUPPORTED, _TORCH_VERSION +from ..._libnatten import HAS_LIBNATTEN +from ...context import is_flex_compile_allowed, is_flex_compile_backprop_allowed +from ...utils.checks import fmha_tensor_checks, log_or_raise_error, na_tensor_checks +from ...utils.device import get_device_cc, is_cpu, is_cuda, is_rocm +from ...utils.dtype import is_fp8 + +### Blackwell FMHA/FNA + + +def can_run_cutlass_blackwell_fmha( + query: Tensor, + key: Tensor, + value: Tensor, + is_causal: bool, + is_varlen: bool, + raise_error: bool = False, +) -> bool: + target_fn = functools.partial(log_or_raise_error, raise_error=raise_error) + + if not HAS_LIBNATTEN: + target_fn("Can't run Blackwell FMHA; NATTEN was not built with libnatten.") + return False + + if not fmha_tensor_checks( + query, + key, + value, + must_match_head_dims=True, + supports_gqa_mqa=True, + raise_error=raise_error, + backend_name="Blackwell FMHA", + ): + return False + + if query.dim() != 4: + target_fn( + f"Blackwell FMHA expects rank-4 input tensors, got {query.shape=}.", + exception=ValueError, + ) + return False + + if not is_cuda(query.device): + target_fn("Can't run Blackwell FMHA; not a CUDA tensor.") + return False + + device_cc = get_device_cc(query.device) + + if device_cc not in [100, 103]: + target_fn( + "Can't run Blackwell FMHA; tensor was on CUDA device with " + f"compute capability {device_cc}, expected 100 or 103." + ) + return False + + head_dim = query.shape[-1] + head_dim_v = value.shape[-1] + + if head_dim != head_dim_v: + target_fn( + "Can't run Blackwell FMHA; it does not support different head dims for QK and V, " + f"got {head_dim=}, {head_dim_v=}.", + exception=ValueError, + ) + return False + + if query.dtype not in [ + torch.float16, + torch.bfloat16, + torch.float8_e4m3fn, + torch.float8_e5m2, + ]: + target_fn( + "Can't run Blackwell FMHA; it only supports FP16, BF16, FP8-E4M3 and FP8-E5M2.", + exception=ValueError, + ) + return False + + requires_grad = query.requires_grad or key.requires_grad or value.requires_grad + if requires_grad and is_fp8(query.dtype): + target_fn( + "Blackwell FMHA does not support FP8 backward pass, but " + f"got {requires_grad=}, {query.dtype=}.", + exception=ValueError, + ) + return False + + if head_dim > 128: + target_fn( + f"Can't run Blackwell FMHA; maximum supported head dim is 128, got {head_dim}.", + exception=NotImplementedError, + ) + return False + + if is_fp8(query.dtype): + if head_dim < 16 or head_dim % 16 != 0: + target_fn( + "Can't run Blackwell FMHA; FP8 requires head dims that are multiples of 16 " + f"(minimum 16), got {head_dim}.", + exception=NotImplementedError, + ) + return False + else: + if head_dim < 8 or head_dim % 8 != 0: + target_fn( + "Can't run Blackwell FMHA; FP16 and BF16 require head dims that are multiples " + f"of 8 (minimum 8), got {head_dim}.", + exception=NotImplementedError, + ) + return False + + return True + + +def can_run_cutlass_blackwell_fna( + query: Tensor, key: Tensor, value: Tensor, raise_error: bool = False +) -> bool: + target_fn = functools.partial(log_or_raise_error, raise_error=raise_error) + + if not HAS_LIBNATTEN: + target_fn("Can't run Blackwell FNA; NATTEN was not built with libnatten.") + return False + + if not na_tensor_checks( + query, + key, + value, + must_match_head_dims=True, + supports_gqa_mqa=True, + raise_error=raise_error, + backend_name="Blackwell FNA", + ): + return False + + if query.dim() not in [4, 5, 6]: + target_fn( + "Blackwell FNA expects 4-D, 5-D, or 6-D tensors as inputs (corresponding to NA1D, NA2D, and NA3D), " + f"got {query.dim()=}.", + exception=ValueError, + ) + return False + + if not is_cuda(query.device): + target_fn("Can't run Blackwell FNA; not a CUDA tensor.") + return False + + device_cc = get_device_cc(query.device) + + if device_cc not in [100, 103]: + target_fn( + "Can't run Blackwell FNA; tensor was on CUDA device with " + f"compute capability {device_cc}, expected 100 or 103." + ) + return False + + requires_grad = query.requires_grad or key.requires_grad or value.requires_grad + if requires_grad and torch.are_deterministic_algorithms_enabled(): + target_fn( + "Can't run Blackwell FNA; its backprop does not have a deterministic mode, but " + "PyTorch's deterministic mode was enabled.", + exception=NotImplementedError, + ) + return False + + head_dim = query.shape[-1] + head_dim_v = value.shape[-1] + + if head_dim != head_dim_v: + target_fn( + "Can't run Blackwell FNA; it does not support different head dims for QK and V, " + f"got {head_dim=}, {head_dim_v=}.", + exception=ValueError, + ) + return False + + if query.dtype not in [ + torch.float16, + torch.bfloat16, + torch.float8_e4m3fn, + torch.float8_e5m2, + ]: + target_fn( + "Can't run Blackwell FNA; it only supports FP16, BF16, FP8-E4M3 and FP8-E5M2.", + exception=ValueError, + ) + return False + + if requires_grad and is_fp8(query.dtype): + target_fn( + "Blackwell FNA does not support FP8 backward pass, but " + f"got {requires_grad=}, {query.dtype=}.", + exception=ValueError, + ) + return False + + if head_dim > 128: + target_fn( + f"Can't run Blackwell FNA; maximum supported head dim is 128, got {head_dim}.", + exception=NotImplementedError, + ) + return False + + if is_fp8(query.dtype): + if head_dim < 16 or head_dim % 16 != 0: + target_fn( + "Can't run Blackwell FNA; FP8 requires head dims that are multiples of 16 " + f"(minimum 16), got {head_dim}.", + exception=NotImplementedError, + ) + return False + else: + if head_dim < 8 or head_dim % 8 != 0: + target_fn( + "Can't run Blackwell FNA; FP16 and BF16 require head dims that are multiples " + f"of 8 (minimum 8), got {head_dim}.", + exception=NotImplementedError, + ) + return False + + return True + + +### Hopper FMHA/FNA + + +def can_run_cutlass_hopper_fmha( + query: Tensor, + key: Tensor, + value: Tensor, + is_causal: bool, + is_varlen: bool, + raise_error: bool = False, +) -> bool: + target_fn = functools.partial(log_or_raise_error, raise_error=raise_error) + + if not HAS_LIBNATTEN: + target_fn("Can't run Hopper FMHA; NATTEN was not built with libnatten.") + return False + + if not fmha_tensor_checks( + query, + key, + value, + must_match_head_dims=True, + supports_gqa_mqa=True, # NOTE: supports GQA in only by virtue of repeating heads manually + raise_error=raise_error, + backend_name="Hopper FMHA", + ): + return False + + if query.dim() != 4: + target_fn( + f"Hopper FMHA expects rank-4 input tensors, got {query.shape=}.", + exception=ValueError, + ) + return False + + if not is_cuda(query.device): + target_fn("Can't run Hopper FMHA; not a CUDA tensor.") + return False + + device_cc = get_device_cc(query.device) + + if device_cc != 90: + target_fn( + "Can't run Hopper FMHA; tensor was on CUDA device with " + f"compute capability {device_cc}, expected 90." + ) + return False + + head_dim = query.shape[-1] + head_dim_v = value.shape[-1] + + if head_dim != head_dim_v: + target_fn( + "Can't run Hopper FMHA; it does not support different head dims for QK and V, " + f"got {head_dim=}, {head_dim_v=}.", + exception=ValueError, + ) + return False + + requires_grad = query.requires_grad or key.requires_grad or value.requires_grad + if requires_grad and head_dim not in [32, 64, 128]: + target_fn( + f"Can't run Hopper FMHA; it does not support backpropagation for {head_dim=} yet; " + "only head dims 32, 64, and 128 are allowed.", + exception=NotImplementedError, + ) + return False + + if requires_grad and torch.are_deterministic_algorithms_enabled(): + target_fn( + "Can't run Hopper FMHA; its backprop does not have a deterministic mode, but " + "PyTorch's deterministic mode was enabled.", + exception=NotImplementedError, + ) + return False + + if query.dtype not in [torch.float16, torch.bfloat16]: + target_fn( + "Can't run Hopper FMHA; it only supports FP16 and BF16 for now.", + exception=ValueError, + ) + return False + + if head_dim not in [32, 64, 128, 256]: + target_fn( + "Can't run Hopper FMHA; it only supports head dims 32, 64, 128, and 256 for now.", + exception=NotImplementedError, + ) + return False + + return True + + +def can_run_cutlass_hopper_fna( + query: Tensor, key: Tensor, value: Tensor, raise_error: bool = False +) -> bool: + target_fn = functools.partial(log_or_raise_error, raise_error=raise_error) + + if not HAS_LIBNATTEN: + target_fn("Can't run Hopper FNA; NATTEN was not built with libnatten.") + return False + + if not na_tensor_checks( + query, + key, + value, + must_match_head_dims=True, + supports_gqa_mqa=True, # NOTE: supports GQA in only by virtue of repeating heads manually + raise_error=raise_error, + backend_name="Hopper FNA", + ): + return False + + if query.dim() not in [4, 5, 6]: + target_fn( + "Hopper FNA expects 4-D, 5-D, or 6-D tensors as inputs (corresponding to NA1D, NA2D, and NA3D), " + f"got {query.dim()=}.", + exception=ValueError, + ) + return False + + if not is_cuda(query.device): + target_fn("Can't run Hopper FNA; not a CUDA tensor.") + return False + + device_cc = get_device_cc(query.device) + + if device_cc != 90: + target_fn( + "Can't run Hopper FNA; tensor was on CUDA device with " + f"compute capability {device_cc}, expected 90." + ) + return False + + head_dim = query.shape[-1] + head_dim_v = value.shape[-1] + + if head_dim != head_dim_v: + target_fn( + "Can't run Hopper FNA; it does not support different head dims for QK and V, " + f"got {head_dim=}, {head_dim_v=}.", + exception=ValueError, + ) + return False + + requires_grad = query.requires_grad or key.requires_grad or value.requires_grad + if requires_grad and head_dim not in [32, 64, 128]: + target_fn( + f"Can't run Hopper FNA; it does not support backpropagation for {head_dim=} yet; " + "only head dims 32, 64, and 128 are allowed.", + exception=NotImplementedError, + ) + return False + + if requires_grad and torch.are_deterministic_algorithms_enabled(): + target_fn( + "Can't run Hopper FNA; its backprop does not have a deterministic mode, but " + "PyTorch's deterministic mode was enabled.", + exception=NotImplementedError, + ) + return False + + if query.dtype not in [torch.float16, torch.bfloat16]: + target_fn( + "Can't run Hopper FNA; it only supports FP16 and BF16 for now.", + exception=ValueError, + ) + return False + + if head_dim not in [32, 64, 128, 256]: + target_fn( + "Can't run Hopper FNA; it only supports head dims 32, 64, 128, and 256 for now.", + exception=NotImplementedError, + ) + return False + + return True + + +### CUTLASS FMHA/FNA + + +def can_run_cutlass_fmha( + query: Tensor, + key: Tensor, + value: Tensor, + is_causal: bool, + is_varlen: bool, + raise_error: bool = False, +) -> bool: + target_fn = functools.partial(log_or_raise_error, raise_error=raise_error) + + if not HAS_LIBNATTEN: + target_fn("Can't run CUTLASS FMHA; NATTEN was not built with libnatten.") + return False + + if not fmha_tensor_checks( + query, + key, + value, + must_match_head_dims=False, + supports_gqa_mqa=True, # NOTE: supports GQA in only by virtue of repeating heads manually + raise_error=raise_error, + backend_name="CUTLASS FMHA", + ): + return False + + if query.dim() != 4: + target_fn( + f"FMHA expects rank-4 input tensors, got {query.shape=}.", + exception=ValueError, + ) + return False + + if not is_cuda(query.device): + target_fn("Can't run CUTLASS FMHA; not a CUDA tensor.") + return False + + device_cc = get_device_cc(query.device) + + if device_cc < 60: + target_fn( + "CUTLASS FMHA only supports CUDA devices with compute capability 60 or higher, " + f"got {device_cc}." + ) + return False + + head_dim = query.shape[-1] + head_dim_v = value.shape[-1] + + if query.dtype not in [torch.float32, torch.float16, torch.bfloat16]: + target_fn( + "Can't run CUTLASS FMHA; it only supports FP32, FP16, and BF16.", + exception=ValueError, + ) + return False + + if head_dim % 8 != 0: + target_fn( + "Can't run CUTLASS FMHA; it only supports head dims that are multiples of 8.", + exception=ValueError, + ) + return False + + if head_dim_v % 8 != 0: + target_fn( + "Can't run CUTLASS FMHA; it only supports head dims that are multiples of 8, " + f"got {head_dim_v=}.", + exception=ValueError, + ) + return False + + if max(head_dim, head_dim_v) > 2**16: + target_fn( + f"Can't run CUTLASS FMHA; it supports max head dim of {2**16}, " + f"got {head_dim=}, {head_dim_v=}.", + exception=ValueError, + ) + return False + + return True + + +def can_run_cutlass_fna( + query: Tensor, key: Tensor, value: Tensor, raise_error: bool = False +) -> bool: + target_fn = functools.partial(log_or_raise_error, raise_error=raise_error) + + if not HAS_LIBNATTEN: + target_fn("Can't run CUTLASS FNA; NATTEN was not built with libnatten.") + return False + + if not na_tensor_checks( + query, + key, + value, + must_match_head_dims=False, + supports_gqa_mqa=True, # NOTE: supports GQA in only by virtue of repeating heads manually + raise_error=raise_error, + backend_name="CUTLASS FNA", + ): + return False + + if query.dim() not in [4, 5, 6]: + target_fn( + "CUTLASS FNA expects 4-D, 5-D, or 6-D tensors as inputs (corresponding to NA1D, NA2D, and NA3D), " + f"got {query.dim()=}.", + exception=ValueError, + ) + return False + + if not is_cuda(query.device): + target_fn("Can't run CUTLASS FNA; not a CUDA tensor.") + return False + + device_cc = get_device_cc(query.device) + + if device_cc < 60: + target_fn( + "CUTLASS FNA only supports CUDA devices with compute capability 60 or higher, " + f"got {device_cc}." + ) + return False + + head_dim = query.shape[-1] + head_dim_v = value.shape[-1] + + if query.dtype not in [torch.float32, torch.float16, torch.bfloat16]: + target_fn( + "Can't run CUTLASS FNA; it only supports FP32, FP16, and BF16.", + exception=ValueError, + ) + return False + + if head_dim % 8 != 0: + target_fn( + "Can't run CUTLASS FNA; it only supports head dims that are multiples of 8, " + f"got {head_dim=}.", + exception=ValueError, + ) + return False + + if head_dim_v % 8 != 0: + target_fn( + "Can't run CUTLASS FNA; it only supports head dims that are multiples of 8, " + f"got {head_dim_v=}.", + exception=ValueError, + ) + return False + + if max(head_dim, head_dim_v) > 2**16: + target_fn( + f"Can't run CUTLASS FNA; it supports max head dim of {2**16}, " + f"got {head_dim=}, {head_dim_v=}.", + exception=ValueError, + ) + return False + + return True + + +### Flex FMHA/FNA + +_FLEX_SUPPORTED = _TORCH_VERSION >= [2, 7] +_FLEX_COMPILE_SUPPORTED = _TORCH_VERSION >= [2, 7] and _IS_TORCH_COMPILE_SUPPORTED + + +def can_run_flex_attention( + query: Tensor, + key: Tensor, + value: Tensor, + torch_compile: bool, + is_causal: bool = False, + is_varlen: bool = False, + raise_error: bool = False, +) -> bool: + target_fn = functools.partial(log_or_raise_error, raise_error=raise_error) + + if is_causal: + target_fn("Flex FMHA doesn't support causal mask yet.") + return False + + if is_varlen: + target_fn("Flex FMHA doesn't support variable length inputs (varlen).") + return False + + if not _FLEX_SUPPORTED: + target_fn("Can't run NATTEN with Flex Attention with torch < 2.7.") + return False + + if torch_compile and not _FLEX_COMPILE_SUPPORTED: + target_fn("Can't run NATTEN with Flex Attention (compiled).)") + return False + + if torch_compile and not is_flex_compile_allowed(): + target_fn( + "NATTEN does not allow compiling Flex Attention. This is because we cannot verify " + "Flex's correctness in all scenarios through NATTEN's tests. You can choose to override " + "this, though it is discouraged, as it may affect your results significantly, " + "by doing:\n" + " from ... import allow_flex_compile\n" + " allow_flex_compile()\n" + ) + return False + + requires_grad = query.requires_grad or key.requires_grad or value.requires_grad + if torch_compile and requires_grad and not is_flex_compile_backprop_allowed(): + target_fn( + "NATTEN does not allow compiling Flex Attention for backpropagation " + "({q,k,v}.requires_grad=True). This is because we cannot verify Flex's correctness " + "in all scenarios through NATTEN's tests. You can choose to override this, though " + "it is HIGHLY discouraged, as it may affect the results of your training significantly, " + "by doing:\n" + " from ... import allow_flex_compile_backprop\n" + " allow_flex_compile_backprop()\n" + ) + return False + + # TODO: can we just have different checks for FMHA vs FNA, like the rest of the backends? + if query.dim() == 4 and key.dim() == 4 and query.shape[1] != key.shape[1]: + supported = fmha_tensor_checks( + query, + key, + value, + must_match_head_dims=True, + supports_gqa_mqa=True, + raise_error=raise_error, + backend_name="Flex FMHA", + ) + else: + supported = na_tensor_checks( + query, + key, + value, + must_match_head_dims=True, + supports_gqa_mqa=True, + raise_error=raise_error, + backend_name="Flex FMHA/FNA", + ) + if not supported: + return False + + if query.dim() not in [4, 5, 6]: + target_fn( + "Flex backend expects 4-D, 5-D, or 6-D tensors as inputs (corresponding to FMHA/NA1D, " + f"NA2D, and NA3D), got {query.dim()=}.", + exception=ValueError, + ) + return False + + if not is_cuda(query.device): + if not is_cpu(query.device) and not is_rocm(query.device): + target_fn( + "Can't run Flex Attention; tensor is not on a CUDA, ROCm, or CPU device: " + f"{query.device.type}" + ) + + return False + # TODO: check if ROCm device supports torch.compile/triton? + + else: + device_cc = get_device_cc(query.device) + + if device_cc < 70: + target_fn( + "Flex Attention (compiled) only supports CUDA devices with compute capability " + f"70 or higher, got {device_cc}." + ) + return False + + head_dim = query.shape[-1] + head_dim_v = value.shape[-1] + + if head_dim != head_dim_v: + target_fn( + "Can't run NATTEN with Flex Attention; we don't support different head dims for QK and " + f"V in this backend yet, got {head_dim=}, {head_dim_v=}.", + exception=ValueError, + ) + return False + + if not torch_compile and query.dtype not in [ + torch.float32, + torch.float16, + torch.bfloat16, + ]: + target_fn( + "Can't run NATTEN with Flex Attention; we only support FP32, FP16, and BF16 for now.", + exception=ValueError, + ) + return False + + if torch_compile and query.dtype not in [torch.float16, torch.bfloat16]: + target_fn( + "Can't run NATTEN with Flex Attention (compiled); we only support FP32, FP16, and BF16 for now.", + exception=ValueError, + ) + return False + + if torch_compile and ( + head_dim < 32 or head_dim > 512 or not math.log2(head_dim).is_integer() + ): + target_fn( + "Can't run NATTEN with Flex Attention (compiled); we only allow 32 <= head_dim <= 512 " + f"and only powers of two, got {head_dim}.", + exception=ValueError, + ) + return False + + if not torch_compile and ( + head_dim < 8 or head_dim > 512 or not math.log2(head_dim).is_integer() + ): + target_fn( + "Can't run NATTEN with Flex Attention (not compiled); we only allow 8 <= head_dim <= 512 " + f"and only powers of two, got {head_dim}.", + exception=ValueError, + ) + return False + + return True diff --git a/natten/torch-ext/natten/backends/configs/cutlass/__init__.py b/natten/torch-ext/natten/backends/configs/cutlass/__init__.py new file mode 100644 index 00000000..7ac0bb72 --- /dev/null +++ b/natten/torch-ext/natten/backends/configs/cutlass/__init__.py @@ -0,0 +1,430 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +from typing import List, Optional, Tuple + +import torch +from torch import Tensor + +# FNA/FMHA forward supports 64x64 and 32x128 GEMM configs in all +# use cases. Some architectures (SM80 and SM90 )have more shared +# memory so they can handle 64x128 GEMMs. + +from ....backends.configs.cutlass.fna_backward_128x128 import ( + _FNA_BACKWARD_128x128_TILE_SIZES, +) +from ....backends.configs.cutlass.fna_backward_128x64 import ( + _FNA_BACKWARD_128x64_TILE_SIZES, +) +from ....backends.configs.cutlass.fna_backward_64x64 import ( + _FNA_BACKWARD_64x64_TILE_SIZES, +) + +# FNA/FMHA backward supports 64x64 GEMM configs in all +# use cases. Some architectures have more shared memory +# so they can handle 128x64 or 128x128 GEMMs, but that +# is also dependent on the GEMM K. + +from ....backends.configs.cutlass.fna_forward_32x128 import ( + _FNA_FORWARD_32x128_TILE_SIZES, +) +from ....backends.configs.cutlass.fna_forward_64x128 import ( + _FNA_FORWARD_64x128_TILE_SIZES, +) +from ....backends.configs.cutlass.fna_forward_64x64 import ( + _FNA_FORWARD_64x64_TILE_SIZES, +) +from ....types import ( + CutlassFmhaBackwardConfigType, + CutlassFmhaForwardConfigType, + CutlassFnaBackwardConfigType, + CutlassFnaForwardConfigType, + DimensionType, +) +from ....utils.checks import check_tile_shape +from ....utils.device import get_device_cc, is_cuda + + +def _get_default_tile_shapes_forward( + na_dim: int, +) -> CutlassFnaForwardConfigType: + assert na_dim in [1, 2, 3] + + if na_dim == 1: + return ((64,), (64,)) + if na_dim == 2: + return ((8, 8), (8, 8)) + if na_dim == 3: + return ((4, 4, 4), (4, 4, 4)) + + raise NotImplementedError() + + +def get_all_tile_shapes_forward( + input_tensor: Tensor, +) -> List[CutlassFnaForwardConfigType]: + assert input_tensor.dim() in [4, 5, 6] + na_dim = input_tensor.dim() - 3 # batch, heads, head_dim + device = input_tensor.device + + if not is_cuda(device): + return [] + + # SM80 and SM90 have more shared memory than SM86 and SM89 + # and their tensor core GEMMs can therefore target larger + # tile shapes. + # SM80 and 86 have been tested, but I don't have an SM89. + # However, I suspect SM89 is to SM90 what SM86 was to SM80 + # in terms of shared memory (and only that). + # Better to disable the larger tile configs for SM89 as well + # as 86 until we can test it. + if get_device_cc(device) in [86, 89]: + return ( + _FNA_FORWARD_32x128_TILE_SIZES[na_dim] + + _FNA_FORWARD_64x64_TILE_SIZES[na_dim] + ) + + return ( + _FNA_FORWARD_32x128_TILE_SIZES[na_dim] + + _FNA_FORWARD_64x64_TILE_SIZES[na_dim] + + _FNA_FORWARD_64x128_TILE_SIZES[na_dim] + ) + + +# For FMHA +def get_all_tile_sizes_forward( + input_tensor: Tensor, +) -> List[CutlassFmhaForwardConfigType]: + if input_tensor.dim() != 4: + raise ValueError("Only 4-D tensors are supported in FMHA.") + + tile_shapes = get_all_tile_shapes_forward(input_tensor) + assert all(len(q_t) == len(kv_t) == 1 for q_t, kv_t in tile_shapes) + + tile_sizes = [(q_t[0], kv_t[0]) for q_t, kv_t in tile_shapes] + + return tile_sizes + + +def get_default_forward_config( + input_tensor: Tensor, dilation: Optional[DimensionType] = None +) -> CutlassFnaForwardConfigType: + assert input_tensor.dim() in [4, 5, 6] + na_dim = input_tensor.dim() - 3 # batch, heads, head_dim + return _get_default_tile_shapes_forward(na_dim) + + +get_all_forward_configs = get_all_tile_shapes_forward +get_all_fmha_forward_configs = get_all_tile_sizes_forward + + +def check_cutlass_fna_forward_config( + input_tensor: Tensor, + dilation: Optional[DimensionType] = None, + q_tile_shape: Optional[DimensionType] = None, + kv_tile_shape: Optional[DimensionType] = None, +) -> CutlassFnaForwardConfigType: + assert input_tensor.dim() in [4, 5, 6] + na_dim = input_tensor.dim() - 3 # batch, heads, head_dim + + if (q_tile_shape is None) ^ (kv_tile_shape is None): + raise ValueError( + "Please specify both q_tile_shape and kv_tile_shape, or neither one. " + f"Got {q_tile_shape=}, {kv_tile_shape=}." + ) + + if q_tile_shape is None and kv_tile_shape is None: + return get_default_forward_config(input_tensor=input_tensor, dilation=dilation) + + q_tile_shape = check_tile_shape(q_tile_shape) + kv_tile_shape = check_tile_shape(kv_tile_shape) + + tile_shapes = get_all_tile_shapes_forward(input_tensor=input_tensor) + + for q_t, kv_t in tile_shapes: + if q_t == q_tile_shape and kv_t == kv_tile_shape: + return (q_t, kv_t) # type: ignore + + # Fail and make suggestions + device_cc = get_device_cc(input_tensor.device) + MAX_EXAMPLES = 3 + examples = "" + for i, (q_t, kv_t) in enumerate(tile_shapes): + examples += f"\n q_tile_shape={q_t}, kv_tile_shape={kv_t}" + if i > MAX_EXAMPLES: + break + + raise ValueError( + f"Invalid configuration for CUTLASS FNA-{na_dim}D. " + f"Q tile shape {q_tile_shape} and KV tile shape {kv_tile_shape} " + f"are not among the {len(tile_shapes)} configurations implementable " + f"with CUTLASS 2.X FNA for SM{device_cc}. " + "Try selecting a combination from: \n" + " natten.get_configs_for_cutlass_fna(q, k, v)" + "\n" + "Here's a few examples of available combinations for your use case:\n" + f"{examples}" + ) + + +def check_cutlass_fmha_forward_config( + input_tensor: Tensor, + q_tile_size: Optional[int] = None, + kv_tile_size: Optional[int] = None, +) -> CutlassFmhaForwardConfigType: + assert input_tensor.dim() == 4 + + if (q_tile_size is None) ^ (kv_tile_size is None): + raise ValueError( + "Please specify both q_tile_size and kv_tile_size, or neither one. " + f"Got {q_tile_size=}, {kv_tile_size=}." + ) + + if q_tile_size is None and kv_tile_size is None: + q_tile_shape, kv_tile_shape = get_default_forward_config( + input_tensor=input_tensor + ) + assert len(q_tile_shape) == len(kv_tile_shape) == 1 + return (q_tile_shape[0], kv_tile_shape[0]) + + tile_sizes = get_all_tile_sizes_forward(input_tensor=input_tensor) + + for q_t, kv_t in tile_sizes: + if q_t == q_tile_size and kv_t == kv_tile_size: + return (q_t, kv_t) + + # Fail and make suggestions + device_cc = get_device_cc(input_tensor.device) + MAX_EXAMPLES = 3 + examples = "" + for i, (q_t, kv_t) in enumerate(tile_sizes): + examples += f"\n q_tile_size={q_t}, kv_tile_size={kv_t}" + if i > MAX_EXAMPLES: + break + + raise ValueError( + f"Invalid configuration for CUTLASS FMHA. " + f"Q tile size {q_tile_size} and KV tile size {kv_tile_size} " + f"are not among the {len(tile_sizes)} configurations implementable " + f"with CUTLASS 2.X FMHA for SM{device_cc}. " + "Try selecting a combination from: \n" + " natten.get_configs_for_cutlass_fmha(q, k, v)" + "\n" + "Here's a few examples of available combinations for your use case:\n" + f"{examples}" + ) + + +###### Backward + + +def _get_default_tile_shapes_backward( + na_dim: int, +) -> Tuple[DimensionType, DimensionType]: + assert na_dim in [1, 2, 3] + + if na_dim == 1: + return ((64,), (64,)) + if na_dim == 2: + return ((8, 8), (8, 8)) + if na_dim == 3: + return ((4, 4, 4), (4, 4, 4)) + + raise NotImplementedError() + + +def get_all_tile_shapes_backward( + input_tensor: Tensor, +) -> List[CutlassFnaForwardConfigType]: + assert input_tensor.dim() in [4, 5, 6] + na_dim = input_tensor.dim() - 3 # batch, heads, head_dim + device = input_tensor.device + dtype = input_tensor.dtype + dim_per_head = input_tensor.shape[-1] + + if not is_cuda(device): + return [] + + compute_cap = get_device_cc(device) + + assert dtype in [torch.float32, torch.float16, torch.bfloat16] + + # DC-class cards have extra shmem which allows larger tile sizes + dc_class_arches = [80, 90, 100, 103] + + if dtype == torch.float32 and compute_cap not in dc_class_arches: + return _FNA_BACKWARD_64x64_TILE_SIZES[na_dim] + + elif dtype == torch.float32: + return ( + _FNA_BACKWARD_64x64_TILE_SIZES[na_dim] + + _FNA_BACKWARD_128x64_TILE_SIZES[na_dim] + ) + + if compute_cap == 70: + return ( + _FNA_BACKWARD_64x64_TILE_SIZES[na_dim] + + _FNA_BACKWARD_128x64_TILE_SIZES[na_dim] + ) + + if compute_cap in dc_class_arches and dim_per_head <= 128: + return ( + _FNA_BACKWARD_64x64_TILE_SIZES[na_dim] + + _FNA_BACKWARD_128x128_TILE_SIZES[na_dim] + ) + elif compute_cap in dc_class_arches: + return ( + _FNA_BACKWARD_64x64_TILE_SIZES[na_dim] + + _FNA_BACKWARD_128x64_TILE_SIZES[na_dim] + ) + + return _FNA_BACKWARD_64x64_TILE_SIZES[na_dim] + + +# For FMHA +def get_all_tile_sizes_backward( + input_tensor: Tensor, +) -> List[CutlassFmhaForwardConfigType]: + if input_tensor.dim() != 4: + raise ValueError("Only 4-D tensors are supported in FMHA.") + + tile_shapes = get_all_tile_shapes_backward(input_tensor) + assert all(len(q_t) == len(kv_t) == 1 for q_t, kv_t in tile_shapes) + + tile_sizes = [(q_t[0], kv_t[0]) for q_t, kv_t in tile_shapes] + + return tile_sizes + + +def get_default_backward_config( + input_tensor: Tensor, +) -> CutlassFnaBackwardConfigType: + assert input_tensor.dim() in [4, 5, 6] + na_dim = input_tensor.dim() - 3 # batch, heads, head_dim + + q_tile_shape, kv_tile_shape = _get_default_tile_shapes_backward(na_dim) + return (q_tile_shape, kv_tile_shape) # type: ignore + + +get_all_backward_configs = get_all_tile_shapes_backward + + +get_all_fmha_backward_configs = get_all_tile_sizes_backward + + +def check_cutlass_fna_backward_config( + input_tensor: Tensor, + q_tile_shape: Optional[DimensionType] = None, + kv_tile_shape: Optional[DimensionType] = None, +) -> CutlassFnaBackwardConfigType: + assert input_tensor.dim() in [4, 5, 6] + na_dim = input_tensor.dim() - 3 # batch, heads, head_dim + + if (q_tile_shape is None) ^ (kv_tile_shape is None): + raise ValueError( + "Please specify both q_tile_shape and kv_tile_shape, or neither one. " + f"Got {q_tile_shape=}, {kv_tile_shape=}." + ) + + if q_tile_shape is None and kv_tile_shape is None: + return get_default_backward_config( + input_tensor=input_tensor, + ) + + q_tile_shape = check_tile_shape(q_tile_shape) + kv_tile_shape = check_tile_shape(kv_tile_shape) + + tile_shapes = get_all_tile_shapes_backward(input_tensor) + + for q_t, kv_t in tile_shapes: + if q_t == q_tile_shape and kv_t == kv_tile_shape: + return (q_t, kv_t) # type: ignore + + # Fail and make suggestions + device_cc = get_device_cc(input_tensor.device) + MAX_EXAMPLES = 3 + examples = "" + for i, (q_t, kv_t) in enumerate(tile_shapes): + examples += f"\n q_tile_shape={q_t}, kv_tile_shape={kv_t}" + if i > MAX_EXAMPLES: + break + + raise ValueError( + f"Invalid configuration for CUTLASS FNA-{na_dim}D. " + f"Q tile shape {q_tile_shape} and KV tile shape {kv_tile_shape} " + f"are not among the {len(tile_shapes)} configurations implementable " + f"with CUTLASS 2.X FNA for SM{device_cc}, with input tensor shape " + f"{input_tensor.shape}. Try selecting a combination from: \n" + " natten.get_bwd_configs_for_cutlass_fna(q, k, v)" + "\n" + "Here's a few examples of available combinations for your use case:\n" + f"{examples}" + ) + + +def check_cutlass_fmha_backward_config( + input_tensor: Tensor, + q_tile_size: Optional[int] = None, + kv_tile_size: Optional[int] = None, +) -> CutlassFmhaBackwardConfigType: + assert input_tensor.dim() == 4 + + if (q_tile_size is None) ^ (kv_tile_size is None): + raise ValueError( + "Please specify both q_tile_size and kv_tile_size, or neither one. " + f"Got {q_tile_size=}, {kv_tile_size=}." + ) + + if q_tile_size is None and kv_tile_size is None: + q_tile_shape, kv_tile_shape = get_default_backward_config( + input_tensor=input_tensor + ) + assert len(q_tile_shape) == len(kv_tile_shape) == 1 + q_tile_size, kv_tile_size = q_tile_shape[0], kv_tile_shape[0] + + tile_sizes = get_all_tile_sizes_backward(input_tensor) + + for q_t, kv_t in tile_sizes: + if q_t == q_tile_size and kv_t == kv_tile_size: + return (q_t, kv_t) + + # Fail and make suggestions + device_cc = get_device_cc(input_tensor.device) + MAX_EXAMPLES = 3 + examples = "" + for i, (q_t, kv_t) in enumerate(tile_sizes): + examples += f"\n q_tile_size={q_t}, kv_tile_size={kv_t}" + if i > MAX_EXAMPLES: + break + + raise ValueError( + f"Invalid configuration for CUTLASS FMHA. " + f"Q tile size {q_tile_size} and KV tile size {kv_tile_size} " + f"are not among the {len(tile_sizes)} configurations implementable " + f"with CUTLASS 2.X FNA for SM{device_cc}, with input tensor shape " + f"{input_tensor.shape}. Try selecting a combination from: \n" + " natten.get_bwd_configs_for_cutlass_fmha(q, k, v)" + "\n" + "Here's a few examples of available combinations for your use case:\n" + f"{examples}" + ) diff --git a/natten/torch-ext/natten/backends/configs/cutlass/backward_knobs.py b/natten/torch-ext/natten/backends/configs/cutlass/backward_knobs.py new file mode 100644 index 00000000..da672aaf --- /dev/null +++ b/natten/torch-ext/natten/backends/configs/cutlass/backward_knobs.py @@ -0,0 +1,228 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +# Backward pass "knobs" for CUTLASS FNA/FMHA: +# - kv_splits: number of KV splits for parallelism +# - use_pt_reduction: whether to use PyTorch for delta computation +# +# These are independent of tile shape selection and are validated/defaulted +# in the torch ops (torch_wrappers.py), not in the config selection logic. + +import itertools +import math +from typing import Optional + +from torch import Tensor + +from ....context import ( + is_kv_parallelism_in_fused_na_enabled, + is_memory_usage_strict, + is_memory_usage_unrestricted, +) +from ....types import DimensionType +from ....utils.checks import check_dilation_arg, check_input_size_arg +from ....utils.tuples import ceil_div_int, ceil_div_tuple + + +def _get_max_grid_size_allowed() -> int: + if is_memory_usage_unrestricted(): + return 65535 + if is_memory_usage_strict(): + return 1024 + + return 4096 + + +def get_min_splits(na_dim: int) -> DimensionType: + assert na_dim in [1, 2, 3] + return tuple(1 for _ in range(na_dim)) # type: ignore + + +def get_max_splits( + input_shape: DimensionType, dilation: DimensionType, kv_tile_shape: DimensionType +) -> DimensionType: + extent_per_dilation_group = ceil_div_tuple(input_shape, dilation) + return tuple( + ceil_div_int(x, t) for x, t in zip(extent_per_dilation_group, kv_tile_shape) + ) # type: ignore + + +def _reduce_max_kv_splits( + na_dim: int, + kv_splits: DimensionType, + max_splits: int, +) -> DimensionType: + assert isinstance(kv_splits, tuple) + assert na_dim in [1, 2, 3] + + if na_dim == 1: + assert len(kv_splits) == 1 + return (min(kv_splits[0], max_splits),) + + if na_dim == 2: + assert len(kv_splits) == 2 + splits_x = max(min(max_splits // 2, kv_splits[0]), 1) + splits_y = max(min(max_splits // splits_x, kv_splits[1]), 1) + assert ( + 0 < splits_x * splits_y <= max_splits + ), f"{splits_x=} * {splits_y=} does not fall in range [0, {max_splits}]" + return (splits_x, splits_y) + + if na_dim == 3: + assert len(kv_splits) == 3 + splits_x = max(min(max_splits // 3, kv_splits[0]), 1) + splits_y = max(min(max_splits // splits_x, kv_splits[1]), 1) + splits_z = max(min(max_splits // (splits_x * splits_y), kv_splits[2]), 1) + assert ( + 0 < splits_x * splits_y * splits_z <= max_splits + ), f"{splits_x=} * {splits_y=} * {splits_z=} does not fall in range [0, {max_splits}]" + return (splits_x, splits_y, splits_z) + + raise NotImplementedError() + + +def _get_possible_kv_splits( + min_splits: DimensionType, + max_splits: DimensionType, +): + assert 0 < len(min_splits) == len(max_splits) < 4 + na_dim = len(max_splits) + if na_dim == 1: + return itertools.product( + range(min_splits[0], max_splits[0] + 1), + ) + if na_dim == 2: + assert len(min_splits) == len(max_splits) == 2 + return itertools.product( + range(min_splits[0], max_splits[0] + 1), + range(min_splits[1], max_splits[1] + 1), + ) + if na_dim == 3: + assert len(min_splits) == len(max_splits) == 3 + return itertools.product( + range(min_splits[0], max_splits[0] + 1), + range(min_splits[1], max_splits[1] + 1), + range(min_splits[2], max_splits[2] + 1), + ) + + raise NotImplementedError() + + +def get_default_kv_splits_backward( + input_tensor: Tensor, + kv_tile_shape: DimensionType, + deterministic: bool, + dilation: Optional[DimensionType] = None, + max_seqlen: Optional[DimensionType] = None, +) -> DimensionType: + assert input_tensor.dim() in [4, 5, 6] + na_dim = input_tensor.dim() - 3 # batch, heads, head_dim + dilation = check_dilation_arg(na_dim, dilation) + input_shape: DimensionType = tuple(int(x) for x in input_tensor.shape[1 : na_dim + 1]) # type: ignore + if max_seqlen is not None: + input_shape = check_input_size_arg(na_dim, max_seqlen) + + assert na_dim in [1, 2, 3] + if na_dim == 1: + kv_splits: DimensionType = (1,) + elif na_dim == 2: + kv_splits = (1, 1) + + elif na_dim == 3: + kv_splits = (1, 1, 1) + + if is_kv_parallelism_in_fused_na_enabled() and not deterministic: + kv_splits = get_max_splits( + input_shape, dilation=dilation, kv_tile_shape=kv_tile_shape + ) + total_kv_splits = math.prod(kv_splits) + + batch_size = input_tensor.shape[0] + num_heads = input_tensor.shape[-2] + num_dilation_splits = math.prod(dilation) + max_kv_splits_allowed = max( + 1, + _get_max_grid_size_allowed() + // (batch_size * num_heads * num_dilation_splits), + ) + + if total_kv_splits > max_kv_splits_allowed: + kv_splits = _reduce_max_kv_splits( + na_dim=na_dim, kv_splits=kv_splits, max_splits=max_kv_splits_allowed + ) + + return kv_splits + + +def check_fmha_kv_splits( + kv_splits: Optional[int], + input_tensor: Tensor, + kv_tile_size: int, + deterministic: bool, + max_seqlen: Optional[int] = None, +) -> int: + if kv_splits is not None and isinstance(kv_splits, int): + seqlen_kv = input_tensor.shape[1] if max_seqlen is None else max_seqlen + num_kv_tiles = (seqlen_kv + kv_tile_size - 1) // kv_tile_size + assert num_kv_tiles > 0 + return min(num_kv_tiles, kv_splits) + + if kv_splits is None: + max_seqlen_tuple = None if max_seqlen is None else (max_seqlen,) + default_kv_splits: DimensionType = get_default_kv_splits_backward( + input_tensor=input_tensor, + deterministic=deterministic, + kv_tile_shape=(kv_tile_size,), + max_seqlen=max_seqlen_tuple, + ) + assert len(default_kv_splits) == 1 + return default_kv_splits[0] + + raise ValueError(f"Invalid type {type(kv_splits)} for kv_splits.") + + +def check_fna_kv_splits( + kv_splits: Optional[DimensionType], + input_tensor: Tensor, + kv_tile_shape: DimensionType, + deterministic: bool, + dilation: Optional[DimensionType] = None, +) -> DimensionType: + if kv_splits is not None and isinstance(kv_splits, tuple): + na_dim = input_tensor.dim() - 3 + dilation = check_dilation_arg(na_dim, dilation) + input_shape: DimensionType = tuple(int(x) for x in input_tensor.shape[1 : na_dim + 1]) # type: ignore + max_kv_splits = get_max_splits( + input_shape, dilation=dilation, kv_tile_shape=kv_tile_shape + ) + return tuple(min(s, m) for s, m in zip(kv_splits, max_kv_splits)) # type: ignore + + if kv_splits is None: + return get_default_kv_splits_backward( + deterministic=deterministic, + input_tensor=input_tensor, + kv_tile_shape=kv_tile_shape, + dilation=dilation, + ) + + raise ValueError(f"Invalid type {type(kv_splits)} for kv_splits.") diff --git a/natten/torch-ext/natten/backends/configs/cutlass/fna_backward_128x128.py b/natten/torch-ext/natten/backends/configs/cutlass/fna_backward_128x128.py new file mode 100644 index 00000000..a7175e21 --- /dev/null +++ b/natten/torch-ext/natten/backends/configs/cutlass/fna_backward_128x128.py @@ -0,0 +1,304 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + + +from typing import Dict, List + +from ....types import QKTileShapeType + +# NOTE: we're excluding tile shapes that include 1 just to +# reduce the giant number of configs down to a reasonable +# amount; otherwise autotuning would take more than a few +# seconds per call which is unacceptable. Tile shapes with +# 1s are rarely selected. + +_FNA_BACKWARD_128x128_TILE_SIZES: Dict[int, List[QKTileShapeType]] = { + 1: [ + ((128,), (128,)), + ], + 2: [ + ((64, 2), (64, 2)), + ((64, 2), (32, 4)), + ((64, 2), (16, 8)), + ((64, 2), (8, 16)), + ((64, 2), (4, 32)), + ((64, 2), (2, 64)), + ((32, 4), (64, 2)), + ((32, 4), (32, 4)), + ((32, 4), (16, 8)), + ((32, 4), (8, 16)), + ((32, 4), (4, 32)), + ((32, 4), (2, 64)), + ((16, 8), (64, 2)), + ((16, 8), (32, 4)), + ((16, 8), (16, 8)), + ((16, 8), (8, 16)), + ((16, 8), (4, 32)), + ((16, 8), (2, 64)), + ((8, 16), (64, 2)), + ((8, 16), (32, 4)), + ((8, 16), (16, 8)), + ((8, 16), (8, 16)), + ((8, 16), (4, 32)), + ((8, 16), (2, 64)), + ((4, 32), (64, 2)), + ((4, 32), (32, 4)), + ((4, 32), (16, 8)), + ((4, 32), (8, 16)), + ((4, 32), (4, 32)), + ((4, 32), (2, 64)), + ((2, 64), (64, 2)), + ((2, 64), (32, 4)), + ((2, 64), (16, 8)), + ((2, 64), (8, 16)), + ((2, 64), (4, 32)), + ((2, 64), (2, 64)), + ], + 3: [ + ((32, 2, 2), (32, 2, 2)), + ((32, 2, 2), (16, 4, 2)), + ((32, 2, 2), (16, 2, 4)), + ((32, 2, 2), (8, 8, 2)), + ((32, 2, 2), (8, 4, 4)), + ((32, 2, 2), (8, 2, 8)), + ((32, 2, 2), (4, 16, 2)), + ((32, 2, 2), (4, 8, 4)), + ((32, 2, 2), (4, 4, 8)), + ((32, 2, 2), (4, 2, 16)), + ((32, 2, 2), (2, 32, 2)), + ((32, 2, 2), (2, 16, 4)), + ((32, 2, 2), (2, 8, 8)), + ((32, 2, 2), (2, 4, 16)), + ((32, 2, 2), (2, 2, 32)), + ((16, 4, 2), (32, 2, 2)), + ((16, 4, 2), (16, 4, 2)), + ((16, 4, 2), (16, 2, 4)), + ((16, 4, 2), (8, 8, 2)), + ((16, 4, 2), (8, 4, 4)), + ((16, 4, 2), (8, 2, 8)), + ((16, 4, 2), (4, 16, 2)), + ((16, 4, 2), (4, 8, 4)), + ((16, 4, 2), (4, 4, 8)), + ((16, 4, 2), (4, 2, 16)), + ((16, 4, 2), (2, 32, 2)), + ((16, 4, 2), (2, 16, 4)), + ((16, 4, 2), (2, 8, 8)), + ((16, 4, 2), (2, 4, 16)), + ((16, 4, 2), (2, 2, 32)), + ((16, 2, 4), (32, 2, 2)), + ((16, 2, 4), (16, 4, 2)), + ((16, 2, 4), (16, 2, 4)), + ((16, 2, 4), (8, 8, 2)), + ((16, 2, 4), (8, 4, 4)), + ((16, 2, 4), (8, 2, 8)), + ((16, 2, 4), (4, 16, 2)), + ((16, 2, 4), (4, 8, 4)), + ((16, 2, 4), (4, 4, 8)), + ((16, 2, 4), (4, 2, 16)), + ((16, 2, 4), (2, 32, 2)), + ((16, 2, 4), (2, 16, 4)), + ((16, 2, 4), (2, 8, 8)), + ((16, 2, 4), (2, 4, 16)), + ((16, 2, 4), (2, 2, 32)), + ((8, 8, 2), (32, 2, 2)), + ((8, 8, 2), (16, 4, 2)), + ((8, 8, 2), (16, 2, 4)), + ((8, 8, 2), (8, 8, 2)), + ((8, 8, 2), (8, 4, 4)), + ((8, 8, 2), (8, 2, 8)), + ((8, 8, 2), (4, 16, 2)), + ((8, 8, 2), (4, 8, 4)), + ((8, 8, 2), (4, 4, 8)), + ((8, 8, 2), (4, 2, 16)), + ((8, 8, 2), (2, 32, 2)), + ((8, 8, 2), (2, 16, 4)), + ((8, 8, 2), (2, 8, 8)), + ((8, 8, 2), (2, 4, 16)), + ((8, 8, 2), (2, 2, 32)), + ((8, 4, 4), (32, 2, 2)), + ((8, 4, 4), (16, 4, 2)), + ((8, 4, 4), (16, 2, 4)), + ((8, 4, 4), (8, 8, 2)), + ((8, 4, 4), (8, 4, 4)), + ((8, 4, 4), (8, 2, 8)), + ((8, 4, 4), (4, 16, 2)), + ((8, 4, 4), (4, 8, 4)), + ((8, 4, 4), (4, 4, 8)), + ((8, 4, 4), (4, 2, 16)), + ((8, 4, 4), (2, 32, 2)), + ((8, 4, 4), (2, 16, 4)), + ((8, 4, 4), (2, 8, 8)), + ((8, 4, 4), (2, 4, 16)), + ((8, 4, 4), (2, 2, 32)), + ((8, 2, 8), (32, 2, 2)), + ((8, 2, 8), (16, 4, 2)), + ((8, 2, 8), (16, 2, 4)), + ((8, 2, 8), (8, 8, 2)), + ((8, 2, 8), (8, 4, 4)), + ((8, 2, 8), (8, 2, 8)), + ((8, 2, 8), (4, 16, 2)), + ((8, 2, 8), (4, 8, 4)), + ((8, 2, 8), (4, 4, 8)), + ((8, 2, 8), (4, 2, 16)), + ((8, 2, 8), (2, 32, 2)), + ((8, 2, 8), (2, 16, 4)), + ((8, 2, 8), (2, 8, 8)), + ((8, 2, 8), (2, 4, 16)), + ((8, 2, 8), (2, 2, 32)), + ((4, 16, 2), (32, 2, 2)), + ((4, 16, 2), (16, 4, 2)), + ((4, 16, 2), (16, 2, 4)), + ((4, 16, 2), (8, 8, 2)), + ((4, 16, 2), (8, 4, 4)), + ((4, 16, 2), (8, 2, 8)), + ((4, 16, 2), (4, 16, 2)), + ((4, 16, 2), (4, 8, 4)), + ((4, 16, 2), (4, 4, 8)), + ((4, 16, 2), (4, 2, 16)), + ((4, 16, 2), (2, 32, 2)), + ((4, 16, 2), (2, 16, 4)), + ((4, 16, 2), (2, 8, 8)), + ((4, 16, 2), (2, 4, 16)), + ((4, 16, 2), (2, 2, 32)), + ((4, 8, 4), (32, 2, 2)), + ((4, 8, 4), (16, 4, 2)), + ((4, 8, 4), (16, 2, 4)), + ((4, 8, 4), (8, 8, 2)), + ((4, 8, 4), (8, 4, 4)), + ((4, 8, 4), (8, 2, 8)), + ((4, 8, 4), (4, 16, 2)), + ((4, 8, 4), (4, 8, 4)), + ((4, 8, 4), (4, 4, 8)), + ((4, 8, 4), (4, 2, 16)), + ((4, 8, 4), (2, 32, 2)), + ((4, 8, 4), (2, 16, 4)), + ((4, 8, 4), (2, 8, 8)), + ((4, 8, 4), (2, 4, 16)), + ((4, 8, 4), (2, 2, 32)), + ((4, 4, 8), (32, 2, 2)), + ((4, 4, 8), (16, 4, 2)), + ((4, 4, 8), (16, 2, 4)), + ((4, 4, 8), (8, 8, 2)), + ((4, 4, 8), (8, 4, 4)), + ((4, 4, 8), (8, 2, 8)), + ((4, 4, 8), (4, 16, 2)), + ((4, 4, 8), (4, 8, 4)), + ((4, 4, 8), (4, 4, 8)), + ((4, 4, 8), (4, 2, 16)), + ((4, 4, 8), (2, 32, 2)), + ((4, 4, 8), (2, 16, 4)), + ((4, 4, 8), (2, 8, 8)), + ((4, 4, 8), (2, 4, 16)), + ((4, 4, 8), (2, 2, 32)), + ((4, 2, 16), (32, 2, 2)), + ((4, 2, 16), (16, 4, 2)), + ((4, 2, 16), (16, 2, 4)), + ((4, 2, 16), (8, 8, 2)), + ((4, 2, 16), (8, 4, 4)), + ((4, 2, 16), (8, 2, 8)), + ((4, 2, 16), (4, 16, 2)), + ((4, 2, 16), (4, 8, 4)), + ((4, 2, 16), (4, 4, 8)), + ((4, 2, 16), (4, 2, 16)), + ((4, 2, 16), (2, 32, 2)), + ((4, 2, 16), (2, 16, 4)), + ((4, 2, 16), (2, 8, 8)), + ((4, 2, 16), (2, 4, 16)), + ((4, 2, 16), (2, 2, 32)), + ((2, 32, 2), (32, 2, 2)), + ((2, 32, 2), (16, 4, 2)), + ((2, 32, 2), (16, 2, 4)), + ((2, 32, 2), (8, 8, 2)), + ((2, 32, 2), (8, 4, 4)), + ((2, 32, 2), (8, 2, 8)), + ((2, 32, 2), (4, 16, 2)), + ((2, 32, 2), (4, 8, 4)), + ((2, 32, 2), (4, 4, 8)), + ((2, 32, 2), (4, 2, 16)), + ((2, 32, 2), (2, 32, 2)), + ((2, 32, 2), (2, 16, 4)), + ((2, 32, 2), (2, 8, 8)), + ((2, 32, 2), (2, 4, 16)), + ((2, 32, 2), (2, 2, 32)), + ((2, 16, 4), (32, 2, 2)), + ((2, 16, 4), (16, 4, 2)), + ((2, 16, 4), (16, 2, 4)), + ((2, 16, 4), (8, 8, 2)), + ((2, 16, 4), (8, 4, 4)), + ((2, 16, 4), (8, 2, 8)), + ((2, 16, 4), (4, 16, 2)), + ((2, 16, 4), (4, 8, 4)), + ((2, 16, 4), (4, 4, 8)), + ((2, 16, 4), (4, 2, 16)), + ((2, 16, 4), (2, 32, 2)), + ((2, 16, 4), (2, 16, 4)), + ((2, 16, 4), (2, 8, 8)), + ((2, 16, 4), (2, 4, 16)), + ((2, 16, 4), (2, 2, 32)), + ((2, 8, 8), (32, 2, 2)), + ((2, 8, 8), (16, 4, 2)), + ((2, 8, 8), (16, 2, 4)), + ((2, 8, 8), (8, 8, 2)), + ((2, 8, 8), (8, 4, 4)), + ((2, 8, 8), (8, 2, 8)), + ((2, 8, 8), (4, 16, 2)), + ((2, 8, 8), (4, 8, 4)), + ((2, 8, 8), (4, 4, 8)), + ((2, 8, 8), (4, 2, 16)), + ((2, 8, 8), (2, 32, 2)), + ((2, 8, 8), (2, 16, 4)), + ((2, 8, 8), (2, 8, 8)), + ((2, 8, 8), (2, 4, 16)), + ((2, 8, 8), (2, 2, 32)), + ((2, 4, 16), (32, 2, 2)), + ((2, 4, 16), (16, 4, 2)), + ((2, 4, 16), (16, 2, 4)), + ((2, 4, 16), (8, 8, 2)), + ((2, 4, 16), (8, 4, 4)), + ((2, 4, 16), (8, 2, 8)), + ((2, 4, 16), (4, 16, 2)), + ((2, 4, 16), (4, 8, 4)), + ((2, 4, 16), (4, 4, 8)), + ((2, 4, 16), (4, 2, 16)), + ((2, 4, 16), (2, 32, 2)), + ((2, 4, 16), (2, 16, 4)), + ((2, 4, 16), (2, 8, 8)), + ((2, 4, 16), (2, 4, 16)), + ((2, 4, 16), (2, 2, 32)), + ((2, 2, 32), (32, 2, 2)), + ((2, 2, 32), (16, 4, 2)), + ((2, 2, 32), (16, 2, 4)), + ((2, 2, 32), (8, 8, 2)), + ((2, 2, 32), (8, 4, 4)), + ((2, 2, 32), (8, 2, 8)), + ((2, 2, 32), (4, 16, 2)), + ((2, 2, 32), (4, 8, 4)), + ((2, 2, 32), (4, 4, 8)), + ((2, 2, 32), (4, 2, 16)), + ((2, 2, 32), (2, 32, 2)), + ((2, 2, 32), (2, 16, 4)), + ((2, 2, 32), (2, 8, 8)), + ((2, 2, 32), (2, 4, 16)), + ((2, 2, 32), (2, 2, 32)), + ], +} diff --git a/natten/torch-ext/natten/backends/configs/cutlass/fna_backward_128x64.py b/natten/torch-ext/natten/backends/configs/cutlass/fna_backward_128x64.py new file mode 100644 index 00000000..bfb1e957 --- /dev/null +++ b/natten/torch-ext/natten/backends/configs/cutlass/fna_backward_128x64.py @@ -0,0 +1,223 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + + +from typing import Dict, List + +from ....types import CutlassFnaForwardConfigType + +# NOTE: we're excluding tile shapes that include 1 just to +# reduce the giant number of configs down to a reasonable +# amount; otherwise autotuning would take more than a few +# seconds per call which is unacceptable. Tile shapes with +# 1s are rarely selected. + +_FNA_BACKWARD_128x64_TILE_SIZES: Dict[int, List[CutlassFnaForwardConfigType]] = { + 1: [ + ((128,), (64,)), + ], + 2: [ + ((64, 2), (32, 2)), + ((64, 2), (16, 4)), + ((64, 2), (8, 8)), + ((64, 2), (4, 16)), + ((64, 2), (2, 32)), + ((32, 4), (32, 2)), + ((32, 4), (16, 4)), + ((32, 4), (8, 8)), + ((32, 4), (4, 16)), + ((32, 4), (2, 32)), + ((16, 8), (32, 2)), + ((16, 8), (16, 4)), + ((16, 8), (8, 8)), + ((16, 8), (4, 16)), + ((16, 8), (2, 32)), + ((8, 16), (32, 2)), + ((8, 16), (16, 4)), + ((8, 16), (8, 8)), + ((8, 16), (4, 16)), + ((8, 16), (2, 32)), + ((4, 32), (32, 2)), + ((4, 32), (16, 4)), + ((4, 32), (8, 8)), + ((4, 32), (4, 16)), + ((4, 32), (2, 32)), + ((2, 64), (32, 2)), + ((2, 64), (16, 4)), + ((2, 64), (8, 8)), + ((2, 64), (4, 16)), + ((2, 64), (2, 32)), + ], + 3: [ + ((32, 2, 2), (16, 2, 2)), + ((32, 2, 2), (8, 4, 2)), + ((32, 2, 2), (8, 2, 4)), + ((32, 2, 2), (4, 8, 2)), + ((32, 2, 2), (4, 4, 4)), + ((32, 2, 2), (4, 2, 8)), + ((32, 2, 2), (2, 16, 2)), + ((32, 2, 2), (2, 8, 4)), + ((32, 2, 2), (2, 4, 8)), + ((32, 2, 2), (2, 2, 16)), + ((16, 4, 2), (16, 2, 2)), + ((16, 4, 2), (8, 4, 2)), + ((16, 4, 2), (8, 2, 4)), + ((16, 4, 2), (4, 8, 2)), + ((16, 4, 2), (4, 4, 4)), + ((16, 4, 2), (4, 2, 8)), + ((16, 4, 2), (2, 16, 2)), + ((16, 4, 2), (2, 8, 4)), + ((16, 4, 2), (2, 4, 8)), + ((16, 4, 2), (2, 2, 16)), + ((16, 2, 4), (16, 2, 2)), + ((16, 2, 4), (8, 4, 2)), + ((16, 2, 4), (8, 2, 4)), + ((16, 2, 4), (4, 8, 2)), + ((16, 2, 4), (4, 4, 4)), + ((16, 2, 4), (4, 2, 8)), + ((16, 2, 4), (2, 16, 2)), + ((16, 2, 4), (2, 8, 4)), + ((16, 2, 4), (2, 4, 8)), + ((16, 2, 4), (2, 2, 16)), + ((8, 8, 2), (16, 2, 2)), + ((8, 8, 2), (8, 4, 2)), + ((8, 8, 2), (8, 2, 4)), + ((8, 8, 2), (4, 8, 2)), + ((8, 8, 2), (4, 4, 4)), + ((8, 8, 2), (4, 2, 8)), + ((8, 8, 2), (2, 16, 2)), + ((8, 8, 2), (2, 8, 4)), + ((8, 8, 2), (2, 4, 8)), + ((8, 8, 2), (2, 2, 16)), + ((8, 4, 4), (16, 2, 2)), + ((8, 4, 4), (8, 4, 2)), + ((8, 4, 4), (8, 2, 4)), + ((8, 4, 4), (4, 8, 2)), + ((8, 4, 4), (4, 4, 4)), + ((8, 4, 4), (4, 2, 8)), + ((8, 4, 4), (2, 16, 2)), + ((8, 4, 4), (2, 8, 4)), + ((8, 4, 4), (2, 4, 8)), + ((8, 4, 4), (2, 2, 16)), + ((8, 2, 8), (16, 2, 2)), + ((8, 2, 8), (8, 4, 2)), + ((8, 2, 8), (8, 2, 4)), + ((8, 2, 8), (4, 8, 2)), + ((8, 2, 8), (4, 4, 4)), + ((8, 2, 8), (4, 2, 8)), + ((8, 2, 8), (2, 16, 2)), + ((8, 2, 8), (2, 8, 4)), + ((8, 2, 8), (2, 4, 8)), + ((8, 2, 8), (2, 2, 16)), + ((4, 16, 2), (16, 2, 2)), + ((4, 16, 2), (8, 4, 2)), + ((4, 16, 2), (8, 2, 4)), + ((4, 16, 2), (4, 8, 2)), + ((4, 16, 2), (4, 4, 4)), + ((4, 16, 2), (4, 2, 8)), + ((4, 16, 2), (2, 16, 2)), + ((4, 16, 2), (2, 8, 4)), + ((4, 16, 2), (2, 4, 8)), + ((4, 16, 2), (2, 2, 16)), + ((4, 8, 4), (16, 2, 2)), + ((4, 8, 4), (8, 4, 2)), + ((4, 8, 4), (8, 2, 4)), + ((4, 8, 4), (4, 8, 2)), + ((4, 8, 4), (4, 4, 4)), + ((4, 8, 4), (4, 2, 8)), + ((4, 8, 4), (2, 16, 2)), + ((4, 8, 4), (2, 8, 4)), + ((4, 8, 4), (2, 4, 8)), + ((4, 8, 4), (2, 2, 16)), + ((4, 4, 8), (16, 2, 2)), + ((4, 4, 8), (8, 4, 2)), + ((4, 4, 8), (8, 2, 4)), + ((4, 4, 8), (4, 8, 2)), + ((4, 4, 8), (4, 4, 4)), + ((4, 4, 8), (4, 2, 8)), + ((4, 4, 8), (2, 16, 2)), + ((4, 4, 8), (2, 8, 4)), + ((4, 4, 8), (2, 4, 8)), + ((4, 4, 8), (2, 2, 16)), + ((4, 2, 16), (16, 2, 2)), + ((4, 2, 16), (8, 4, 2)), + ((4, 2, 16), (8, 2, 4)), + ((4, 2, 16), (4, 8, 2)), + ((4, 2, 16), (4, 4, 4)), + ((4, 2, 16), (4, 2, 8)), + ((4, 2, 16), (2, 16, 2)), + ((4, 2, 16), (2, 8, 4)), + ((4, 2, 16), (2, 4, 8)), + ((4, 2, 16), (2, 2, 16)), + ((2, 32, 2), (16, 2, 2)), + ((2, 32, 2), (8, 4, 2)), + ((2, 32, 2), (8, 2, 4)), + ((2, 32, 2), (4, 8, 2)), + ((2, 32, 2), (4, 4, 4)), + ((2, 32, 2), (4, 2, 8)), + ((2, 32, 2), (2, 16, 2)), + ((2, 32, 2), (2, 8, 4)), + ((2, 32, 2), (2, 4, 8)), + ((2, 32, 2), (2, 2, 16)), + ((2, 16, 4), (16, 2, 2)), + ((2, 16, 4), (8, 4, 2)), + ((2, 16, 4), (8, 2, 4)), + ((2, 16, 4), (4, 8, 2)), + ((2, 16, 4), (4, 4, 4)), + ((2, 16, 4), (4, 2, 8)), + ((2, 16, 4), (2, 16, 2)), + ((2, 16, 4), (2, 8, 4)), + ((2, 16, 4), (2, 4, 8)), + ((2, 16, 4), (2, 2, 16)), + ((2, 8, 8), (16, 2, 2)), + ((2, 8, 8), (8, 4, 2)), + ((2, 8, 8), (8, 2, 4)), + ((2, 8, 8), (4, 8, 2)), + ((2, 8, 8), (4, 4, 4)), + ((2, 8, 8), (4, 2, 8)), + ((2, 8, 8), (2, 16, 2)), + ((2, 8, 8), (2, 8, 4)), + ((2, 8, 8), (2, 4, 8)), + ((2, 8, 8), (2, 2, 16)), + ((2, 4, 16), (16, 2, 2)), + ((2, 4, 16), (8, 4, 2)), + ((2, 4, 16), (8, 2, 4)), + ((2, 4, 16), (4, 8, 2)), + ((2, 4, 16), (4, 4, 4)), + ((2, 4, 16), (4, 2, 8)), + ((2, 4, 16), (2, 16, 2)), + ((2, 4, 16), (2, 8, 4)), + ((2, 4, 16), (2, 4, 8)), + ((2, 4, 16), (2, 2, 16)), + ((2, 2, 32), (16, 2, 2)), + ((2, 2, 32), (8, 4, 2)), + ((2, 2, 32), (8, 2, 4)), + ((2, 2, 32), (4, 8, 2)), + ((2, 2, 32), (4, 4, 4)), + ((2, 2, 32), (4, 2, 8)), + ((2, 2, 32), (2, 16, 2)), + ((2, 2, 32), (2, 8, 4)), + ((2, 2, 32), (2, 4, 8)), + ((2, 2, 32), (2, 2, 16)), + ], +} diff --git a/natten/torch-ext/natten/backends/configs/cutlass/fna_backward_64x64.py b/natten/torch-ext/natten/backends/configs/cutlass/fna_backward_64x64.py new file mode 100644 index 00000000..28e9090a --- /dev/null +++ b/natten/torch-ext/natten/backends/configs/cutlass/fna_backward_64x64.py @@ -0,0 +1,168 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + + +from typing import Dict, List + +from ....types import QKTileShapeType + +# NOTE: we're excluding tile shapes that include 1 just to +# reduce the giant number of configs down to a reasonable +# amount; otherwise autotuning would take more than a few +# seconds per call which is unacceptable. Tile shapes with +# 1s are rarely selected. + +_FNA_BACKWARD_64x64_TILE_SIZES: Dict[int, List[QKTileShapeType]] = { + 1: [ + ((64,), (64,)), + ], + 2: [ + ((32, 2), (32, 2)), + ((32, 2), (16, 4)), + ((32, 2), (8, 8)), + ((32, 2), (4, 16)), + ((32, 2), (2, 32)), + ((16, 4), (32, 2)), + ((16, 4), (16, 4)), + ((16, 4), (8, 8)), + ((16, 4), (4, 16)), + ((16, 4), (2, 32)), + ((8, 8), (32, 2)), + ((8, 8), (16, 4)), + ((8, 8), (8, 8)), + ((8, 8), (4, 16)), + ((8, 8), (2, 32)), + ((4, 16), (32, 2)), + ((4, 16), (16, 4)), + ((4, 16), (8, 8)), + ((4, 16), (4, 16)), + ((4, 16), (2, 32)), + ((2, 32), (32, 2)), + ((2, 32), (16, 4)), + ((2, 32), (8, 8)), + ((2, 32), (4, 16)), + ((2, 32), (2, 32)), + ], + 3: [ + ((16, 2, 2), (16, 2, 2)), + ((16, 2, 2), (8, 4, 2)), + ((16, 2, 2), (8, 2, 4)), + ((16, 2, 2), (4, 8, 2)), + ((16, 2, 2), (4, 4, 4)), + ((16, 2, 2), (4, 2, 8)), + ((16, 2, 2), (2, 16, 2)), + ((16, 2, 2), (2, 8, 4)), + ((16, 2, 2), (2, 4, 8)), + ((16, 2, 2), (2, 2, 16)), + ((8, 4, 2), (16, 2, 2)), + ((8, 4, 2), (8, 4, 2)), + ((8, 4, 2), (8, 2, 4)), + ((8, 4, 2), (4, 8, 2)), + ((8, 4, 2), (4, 4, 4)), + ((8, 4, 2), (4, 2, 8)), + ((8, 4, 2), (2, 16, 2)), + ((8, 4, 2), (2, 8, 4)), + ((8, 4, 2), (2, 4, 8)), + ((8, 4, 2), (2, 2, 16)), + ((8, 2, 4), (16, 2, 2)), + ((8, 2, 4), (8, 4, 2)), + ((8, 2, 4), (8, 2, 4)), + ((8, 2, 4), (4, 8, 2)), + ((8, 2, 4), (4, 4, 4)), + ((8, 2, 4), (4, 2, 8)), + ((8, 2, 4), (2, 16, 2)), + ((8, 2, 4), (2, 8, 4)), + ((8, 2, 4), (2, 4, 8)), + ((8, 2, 4), (2, 2, 16)), + ((4, 8, 2), (16, 2, 2)), + ((4, 8, 2), (8, 4, 2)), + ((4, 8, 2), (8, 2, 4)), + ((4, 8, 2), (4, 8, 2)), + ((4, 8, 2), (4, 4, 4)), + ((4, 8, 2), (4, 2, 8)), + ((4, 8, 2), (2, 16, 2)), + ((4, 8, 2), (2, 8, 4)), + ((4, 8, 2), (2, 4, 8)), + ((4, 8, 2), (2, 2, 16)), + ((4, 4, 4), (16, 2, 2)), + ((4, 4, 4), (8, 4, 2)), + ((4, 4, 4), (8, 2, 4)), + ((4, 4, 4), (4, 8, 2)), + ((4, 4, 4), (4, 4, 4)), + ((4, 4, 4), (4, 2, 8)), + ((4, 4, 4), (2, 16, 2)), + ((4, 4, 4), (2, 8, 4)), + ((4, 4, 4), (2, 4, 8)), + ((4, 4, 4), (2, 2, 16)), + ((4, 2, 8), (16, 2, 2)), + ((4, 2, 8), (8, 4, 2)), + ((4, 2, 8), (8, 2, 4)), + ((4, 2, 8), (4, 8, 2)), + ((4, 2, 8), (4, 4, 4)), + ((4, 2, 8), (4, 2, 8)), + ((4, 2, 8), (2, 16, 2)), + ((4, 2, 8), (2, 8, 4)), + ((4, 2, 8), (2, 4, 8)), + ((4, 2, 8), (2, 2, 16)), + ((2, 16, 2), (16, 2, 2)), + ((2, 16, 2), (8, 4, 2)), + ((2, 16, 2), (8, 2, 4)), + ((2, 16, 2), (4, 8, 2)), + ((2, 16, 2), (4, 4, 4)), + ((2, 16, 2), (4, 2, 8)), + ((2, 16, 2), (2, 16, 2)), + ((2, 16, 2), (2, 8, 4)), + ((2, 16, 2), (2, 4, 8)), + ((2, 16, 2), (2, 2, 16)), + ((2, 8, 4), (16, 2, 2)), + ((2, 8, 4), (8, 4, 2)), + ((2, 8, 4), (8, 2, 4)), + ((2, 8, 4), (4, 8, 2)), + ((2, 8, 4), (4, 4, 4)), + ((2, 8, 4), (4, 2, 8)), + ((2, 8, 4), (2, 16, 2)), + ((2, 8, 4), (2, 8, 4)), + ((2, 8, 4), (2, 4, 8)), + ((2, 8, 4), (2, 2, 16)), + ((2, 4, 8), (16, 2, 2)), + ((2, 4, 8), (8, 4, 2)), + ((2, 4, 8), (8, 2, 4)), + ((2, 4, 8), (4, 8, 2)), + ((2, 4, 8), (4, 4, 4)), + ((2, 4, 8), (4, 2, 8)), + ((2, 4, 8), (2, 16, 2)), + ((2, 4, 8), (2, 8, 4)), + ((2, 4, 8), (2, 4, 8)), + ((2, 4, 8), (2, 2, 16)), + ((2, 2, 16), (16, 2, 2)), + ((2, 2, 16), (8, 4, 2)), + ((2, 2, 16), (8, 2, 4)), + ((2, 2, 16), (4, 8, 2)), + ((2, 2, 16), (4, 4, 4)), + ((2, 2, 16), (4, 2, 8)), + ((2, 2, 16), (2, 16, 2)), + ((2, 2, 16), (2, 8, 4)), + ((2, 2, 16), (2, 4, 8)), + ((2, 2, 16), (2, 2, 16)), + ], +} diff --git a/natten/torch-ext/natten/backends/configs/cutlass/fna_forward_32x128.py b/natten/torch-ext/natten/backends/configs/cutlass/fna_forward_32x128.py new file mode 100644 index 00000000..d8e0fd2e --- /dev/null +++ b/natten/torch-ext/natten/backends/configs/cutlass/fna_forward_32x128.py @@ -0,0 +1,90 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + + +from typing import Dict, List + +from ....types import QKTileShapeType + +# TODO: More combinations are possible for +# 2D and 3D (query tile does not have to be smaller +# than KV tile); but that behavior is untested, +# and IIRC was unstable. + +_FNA_FORWARD_32x128_TILE_SIZES: Dict[int, List[QKTileShapeType]] = { + 1: [ + ((32,), (128,)), + ], + 2: [ + ((16, 2), (64, 2)), + ((16, 2), (32, 4)), + ((16, 2), (16, 8)), + ((8, 4), (32, 4)), + ((8, 4), (16, 8)), + ((8, 4), (8, 16)), + ((4, 8), (16, 8)), + ((4, 8), (8, 16)), + ((4, 8), (4, 32)), + ((2, 16), (8, 16)), + ((2, 16), (4, 32)), + ((2, 16), (2, 64)), + ], + 3: [ + ((8, 2, 2), (32, 2, 2)), + ((8, 2, 2), (16, 4, 2)), + ((8, 2, 2), (16, 2, 4)), + ((8, 2, 2), (8, 8, 2)), + ((8, 2, 2), (8, 4, 4)), + ((8, 2, 2), (8, 2, 8)), + ((4, 4, 2), (16, 4, 2)), + ((4, 4, 2), (8, 8, 2)), + ((4, 4, 2), (8, 4, 4)), + ((4, 4, 2), (4, 16, 2)), + ((4, 4, 2), (4, 8, 4)), + ((4, 4, 2), (4, 4, 8)), + ((4, 2, 4), (16, 2, 4)), + ((4, 2, 4), (8, 4, 4)), + ((4, 2, 4), (8, 2, 8)), + ((4, 2, 4), (4, 8, 4)), + ((4, 2, 4), (4, 4, 8)), + ((4, 2, 4), (4, 2, 16)), + ((2, 8, 2), (8, 8, 2)), + ((2, 8, 2), (4, 16, 2)), + ((2, 8, 2), (4, 8, 4)), + ((2, 8, 2), (2, 32, 2)), + ((2, 8, 2), (2, 16, 4)), + ((2, 8, 2), (2, 8, 8)), + ((2, 4, 4), (8, 4, 4)), + ((2, 4, 4), (4, 8, 4)), + ((2, 4, 4), (4, 4, 8)), + ((2, 4, 4), (2, 16, 4)), + ((2, 4, 4), (2, 8, 8)), + ((2, 4, 4), (2, 4, 16)), + ((2, 2, 8), (8, 2, 8)), + ((2, 2, 8), (4, 4, 8)), + ((2, 2, 8), (4, 2, 16)), + ((2, 2, 8), (2, 8, 8)), + ((2, 2, 8), (2, 4, 16)), + ((2, 2, 8), (2, 2, 32)), + ], +} diff --git a/natten/torch-ext/natten/backends/configs/cutlass/fna_forward_64x128.py b/natten/torch-ext/natten/backends/configs/cutlass/fna_forward_64x128.py new file mode 100644 index 00000000..e6309db4 --- /dev/null +++ b/natten/torch-ext/natten/backends/configs/cutlass/fna_forward_64x128.py @@ -0,0 +1,82 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + + +from typing import Dict, List + +from ....types import QKTileShapeType + +# TODO: More combinations are possible for +# 2D and 3D (query tile does not have to be smaller +# than KV tile); but that behavior is untested, +# and IIRC was unstable. + +_FNA_FORWARD_64x128_TILE_SIZES: Dict[int, List[QKTileShapeType]] = { + 1: [ + ((64,), (128,)), + ], + 2: [ + ((32, 2), (64, 2)), + ((32, 2), (32, 4)), + ((16, 4), (32, 4)), + ((16, 4), (16, 8)), + ((8, 8), (16, 8)), + ((8, 8), (8, 16)), + ((4, 16), (8, 16)), + ((4, 16), (4, 32)), + ((2, 32), (4, 32)), + ((2, 32), (2, 64)), + ], + 3: [ + ((16, 2, 2), (32, 2, 2)), + ((16, 2, 2), (16, 4, 2)), + ((16, 2, 2), (16, 2, 4)), + ((8, 4, 2), (16, 4, 2)), + ((8, 4, 2), (8, 8, 2)), + ((8, 4, 2), (8, 4, 4)), + ((8, 2, 4), (16, 2, 4)), + ((8, 2, 4), (8, 4, 4)), + ((8, 2, 4), (8, 2, 8)), + ((4, 8, 2), (8, 8, 2)), + ((4, 8, 2), (4, 16, 2)), + ((4, 8, 2), (4, 8, 4)), + ((4, 4, 4), (8, 4, 4)), + ((4, 4, 4), (4, 8, 4)), + ((4, 4, 4), (4, 4, 8)), + ((4, 2, 8), (8, 2, 8)), + ((4, 2, 8), (4, 4, 8)), + ((4, 2, 8), (4, 2, 16)), + ((2, 16, 2), (4, 16, 2)), + ((2, 16, 2), (2, 32, 2)), + ((2, 16, 2), (2, 16, 4)), + ((2, 8, 4), (4, 8, 4)), + ((2, 8, 4), (2, 16, 4)), + ((2, 8, 4), (2, 8, 8)), + ((2, 4, 8), (4, 4, 8)), + ((2, 4, 8), (2, 8, 8)), + ((2, 4, 8), (2, 4, 16)), + ((2, 2, 16), (4, 2, 16)), + ((2, 2, 16), (2, 4, 16)), + ((2, 2, 16), (2, 2, 32)), + ], +} diff --git a/natten/torch-ext/natten/backends/configs/cutlass/fna_forward_64x64.py b/natten/torch-ext/natten/backends/configs/cutlass/fna_forward_64x64.py new file mode 100644 index 00000000..e70dc716 --- /dev/null +++ b/natten/torch-ext/natten/backends/configs/cutlass/fna_forward_64x64.py @@ -0,0 +1,63 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + + +from typing import Dict, List + +from ....types import QKTileShapeType + +# TODO: More combinations are possible for +# 2D and 3D (query tile does not have to be smaller +# than KV tile); but that behavior is untested, +# and IIRC was unstable. + +# NOTE: we're excluding tile shapes that include 1 just to +# reduce the giant number of configs down to a reasonable +# amount; otherwise autotuning would take more than a few +# seconds per call which is unacceptable. Tile shapes with +# 1s are rarely selected. + +_FNA_FORWARD_64x64_TILE_SIZES: Dict[int, List[QKTileShapeType]] = { + 1: [ + ((64,), (64,)), + ], + 2: [ + ((32, 2), (32, 2)), + ((16, 4), (16, 4)), + ((8, 8), (8, 8)), + ((4, 16), (4, 16)), + ((2, 32), (2, 32)), + ], + 3: [ + ((16, 2, 2), (16, 2, 2)), + ((8, 4, 2), (8, 4, 2)), + ((8, 2, 4), (8, 2, 4)), + ((4, 8, 2), (4, 8, 2)), + ((4, 4, 4), (4, 4, 4)), + ((4, 2, 8), (4, 2, 8)), + ((2, 16, 2), (2, 16, 2)), + ((2, 8, 4), (2, 8, 4)), + ((2, 4, 8), (2, 4, 8)), + ((2, 2, 16), (2, 2, 16)), + ], +} diff --git a/natten/torch-ext/natten/backends/configs/cutlass_blackwell/__init__.py b/natten/torch-ext/natten/backends/configs/cutlass_blackwell/__init__.py new file mode 100644 index 00000000..165fb92f --- /dev/null +++ b/natten/torch-ext/natten/backends/configs/cutlass_blackwell/__init__.py @@ -0,0 +1,391 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +from typing import List, Optional + +import torch # noqa: F401 +from torch import Tensor + +from ....types import ( + CutlassBlackwellFmhaBackwardConfigType, + CutlassBlackwellFmhaForwardConfigType, + CutlassBlackwellFnaBackwardConfigType, + CutlassBlackwellFnaForwardConfigType, + DimensionType, +) +from ....utils.checks import check_tile_shape +from ....utils.device import get_device_cc + +# The current CUTLASS FMHA forward kernel can only do Q tile size 256, KV tile size 128. +# This limits 1D tile shapes to just the one, but for 2-D and 3-D we can have many more shapes, +# only some of which we compile. Adding new ones requires adding them to autogen, regenerating +# the instantiations, and recompiling libnatten. Unlike CUTLASS 2.X FNA, multi-dim tile shapes are +# static in Blackwell FNA, and not dynamic. + +BLACKWELL_FORWARD_TILE_SHAPES = { + 1: [ + ((256,), (128,)), + ], + 2: [ + ((16, 16), (16, 8)), + ((16, 16), (8, 16)), + ((8, 32), (8, 16)), + ((8, 32), (4, 32)), + ], + 3: [ + ((8, 4, 8), (4, 4, 8)), + ((8, 4, 8), (2, 8, 8)), + ((2, 8, 16), (4, 4, 8)), + ((2, 8, 16), (2, 8, 8)), + ((4, 4, 16), (2, 4, 16)), + ((2, 16, 8), (2, 8, 8)), + ((4, 8, 8), (2, 8, 8)), + ], +} + +BLACKWELL_BACKWARD_TILE_SHAPES = { + 1: [ + ((128,), (128,)), + ], + 2: [ + ((16, 8), (16, 8)), + ((16, 8), (8, 16)), + ((8, 16), (16, 8)), + ((8, 16), (8, 16)), + ], + 3: [ + ((4, 4, 8), (4, 4, 8)), + ((4, 4, 8), (2, 8, 8)), + ((1, 8, 16), (4, 4, 8)), + ((2, 8, 8), (4, 4, 8)), + ((1, 8, 16), (2, 8, 8)), + ((2, 4, 16), (2, 4, 16)), + ((4, 2, 16), (2, 4, 16)), + ((4, 4, 8), (2, 4, 16)), + ((2, 8, 8), (2, 8, 8)), + ], +} + + +def _get_default_tile_shapes_forward( + na_dim: int, +) -> CutlassBlackwellFnaForwardConfigType: + assert na_dim in [1, 2, 3] + + if na_dim == 1: + return ((256,), (128,)) + if na_dim == 2: + return ((16, 16), (16, 8)) + if na_dim == 3: + return ((8, 4, 8), (4, 4, 8)) + + raise NotImplementedError() + + +def get_all_forward_configs( + input_tensor: Tensor, +) -> List[CutlassBlackwellFnaForwardConfigType]: + assert input_tensor.dim() in [4, 5, 6] + na_dim = input_tensor.dim() - 3 # batch, heads, head_dim + + device_cc = get_device_cc(input_tensor.device) + if device_cc not in [100, 103]: + return [] + + return BLACKWELL_FORWARD_TILE_SHAPES[na_dim] # type: ignore + + +def get_all_backward_configs( + input_tensor: Tensor, +) -> List[CutlassBlackwellFnaBackwardConfigType]: + assert input_tensor.dim() in [4, 5, 6] + na_dim = input_tensor.dim() - 3 # batch, heads, head_dim + + device_cc = get_device_cc(input_tensor.device) + if device_cc not in [100, 103]: + return [] + + return BLACKWELL_BACKWARD_TILE_SHAPES[na_dim] # type: ignore + + +# For FMHA +def get_all_fmha_forward_configs( + input_tensor: Tensor, +) -> List[CutlassBlackwellFmhaForwardConfigType]: + if input_tensor.dim() != 4: + raise ValueError("Only 4-D tensors are supported in FMHA.") + + tile_shapes = get_all_forward_configs(input_tensor) + assert all(len(q_t) == len(kv_t) == 1 for q_t, kv_t in tile_shapes) + + tile_sizes = [(q_t[0], kv_t[0]) for q_t, kv_t in tile_shapes] + + return tile_sizes + + +def get_all_fmha_backward_configs( + input_tensor: Tensor, +) -> List[CutlassBlackwellFmhaBackwardConfigType]: + if input_tensor.dim() != 4: + raise ValueError("Only 4-D tensors are supported in FMHA.") + + tile_shapes = get_all_backward_configs(input_tensor) + assert all(len(q_t) == len(kv_t) == 1 for q_t, kv_t in tile_shapes) + + tile_sizes = [(q_t[0], kv_t[0]) for q_t, kv_t in tile_shapes] + + return tile_sizes + + +def get_default_forward_tile_shapes( + input_tensor: Tensor, +) -> CutlassBlackwellFnaForwardConfigType: + assert input_tensor.dim() in [4, 5, 6] + na_dim = input_tensor.dim() - 3 # batch, heads, head_dim + + return _get_default_tile_shapes_forward(na_dim) + + +def get_default_forward_tile_sizes( + input_tensor: Tensor, +) -> CutlassBlackwellFmhaForwardConfigType: + if input_tensor.dim() != 4: + raise ValueError("Only 4-D tensors are supported in FMHA.") + + q_t, kv_t = get_default_forward_tile_shapes(input_tensor) + assert len(q_t) == len(kv_t) == 1 + + return (q_t[0], kv_t[0]) + + +def get_default_backward_tile_shapes( + input_tensor: Tensor, +) -> CutlassBlackwellFnaBackwardConfigType: + all_configs = get_all_backward_configs(input_tensor) + + if len(all_configs) < 1: + device_cc = get_device_cc(input_tensor.device) + raise ValueError( + "No configs exist for this use case; Blackwell FMHA/FNA does not support it: " + f"{input_tensor.shape=}, {input_tensor.dtype=}, {device_cc=}." + ) + + return all_configs[0] + + +def get_default_backward_tile_sizes( + input_tensor: Tensor, +) -> CutlassBlackwellFmhaBackwardConfigType: + if input_tensor.dim() != 4: + raise ValueError("Only 4-D tensors are supported in FMHA.") + + q_t, kv_t = get_default_backward_tile_shapes(input_tensor) + assert len(q_t) == len(kv_t) == 1 + + return (q_t[0], kv_t[0]) + + +def check_cutlass_blackwell_fna_forward_config( + input_tensor: Tensor, + q_tile_shape: Optional[DimensionType] = None, + kv_tile_shape: Optional[DimensionType] = None, +) -> CutlassBlackwellFnaForwardConfigType: + assert input_tensor.dim() in [4, 5, 6] + na_dim = input_tensor.dim() - 3 # batch, heads, head_dim + + if (q_tile_shape is None) ^ (kv_tile_shape is None): + raise ValueError( + "Please specify both q_tile_shape and kv_tile_shape, or neither one. " + f"Got {q_tile_shape=}, {kv_tile_shape=}." + ) + + if q_tile_shape is None and kv_tile_shape is None: + return get_default_forward_tile_shapes(input_tensor=input_tensor) + + q_tile_shape = check_tile_shape(q_tile_shape) + kv_tile_shape = check_tile_shape(kv_tile_shape) + + tile_shapes = get_all_forward_configs(input_tensor=input_tensor) + + for q_t, kv_t in tile_shapes: + if q_t == q_tile_shape and kv_t == kv_tile_shape: + return (q_t, kv_t) # type: ignore + + # Fail and make suggestions + MAX_EXAMPLES = 3 + examples = "" + for i, (q_t, kv_t) in enumerate(tile_shapes): + examples += f"\n q_tile_shape={q_t}, kv_tile_shape={kv_t}" + if i > MAX_EXAMPLES: + break + + raise ValueError( + f"Invalid configuration for CUTLASS Blackwell FNA-{na_dim}D. " + f"Q tile shape {q_tile_shape} and KV tile shape {kv_tile_shape} " + f"are not among the {len(tile_shapes)} configurations implementable " + f"with CUTLASS Blackwell FNA. " + "Try selecting a combination from: \n" + " natten.get_configs_for_blackwell_fna(q, k, v)" + "\n" + "Here's a few examples of available combinations for your use case:\n" + f"{examples}" + ) + + +def check_cutlass_blackwell_fmha_forward_config( + input_tensor: Tensor, + q_tile_size: Optional[int] = None, + kv_tile_size: Optional[int] = None, +) -> CutlassBlackwellFmhaForwardConfigType: + assert input_tensor.dim() == 4 + + if (q_tile_size is None) ^ (kv_tile_size is None): + raise ValueError( + "Please specify both q_tile_size and kv_tile_size, or neither one. " + f"Got {q_tile_size=}, {kv_tile_size=}." + ) + + if q_tile_size is None and kv_tile_size is None: + q_tile_size, kv_tile_size = get_default_forward_tile_sizes( + input_tensor=input_tensor + ) + return (q_tile_size, kv_tile_size) + + tile_sizes = get_all_fmha_forward_configs(input_tensor=input_tensor) + + for q_t, kv_t in tile_sizes: + if q_t == q_tile_size and kv_t == kv_tile_size: + return (q_t, kv_t) + + # Fail and make suggestions + MAX_EXAMPLES = 3 + examples = "" + for i, (q_t, kv_t) in enumerate(tile_sizes): + examples += f"\n q_tile_size={q_t}, kv_tile_size={kv_t}" + if i > MAX_EXAMPLES: + break + + raise ValueError( + f"Invalid configuration for CUTLASS Blackwell FMHA. " + f"Q tile size {q_tile_size} and KV tile size {kv_tile_size} " + f"are not among the {len(tile_sizes)} configurations implementable " + f"with CUTLASS Blackwell FMHA. " + "Try selecting a combination from: \n" + " natten.get_configs_for_blackwell_fmha(q, k, v)" + "\n" + "Here's a few examples of available combinations for your use case:\n" + f"{examples}" + ) + + +def check_cutlass_blackwell_fna_backward_config( + input_tensor: Tensor, + q_tile_shape: Optional[DimensionType] = None, + kv_tile_shape: Optional[DimensionType] = None, +) -> CutlassBlackwellFnaBackwardConfigType: + assert input_tensor.dim() in [4, 5, 6] + na_dim = input_tensor.dim() - 3 # batch, heads, head_dim + + if (q_tile_shape is None) ^ (kv_tile_shape is None): + raise ValueError( + "Please specify both q_tile_shape and kv_tile_shape, or neither one. " + f"Got {q_tile_shape=}, {kv_tile_shape=}." + ) + + if q_tile_shape is None and kv_tile_shape is None: + return get_default_backward_tile_shapes(input_tensor=input_tensor) + + q_tile_shape = check_tile_shape(q_tile_shape) + kv_tile_shape = check_tile_shape(kv_tile_shape) + + tile_shapes = get_all_backward_configs(input_tensor=input_tensor) + + for q_t, kv_t in tile_shapes: + if q_t == q_tile_shape and kv_t == kv_tile_shape: + return (q_t, kv_t) # type: ignore + + # Fail and make suggestions + MAX_EXAMPLES = 3 + examples = "" + for i, (q_t, kv_t) in enumerate(tile_shapes): + examples += f"\n q_tile_shape={q_t}, kv_tile_shape={kv_t}" + if i > MAX_EXAMPLES: + break + + raise ValueError( + f"Invalid configuration for CUTLASS Blackwell FNA-{na_dim}D Backward. " + f"Q tile shape {q_tile_shape} and KV tile shape {kv_tile_shape} " + f"are not among the {len(tile_shapes)} configurations implementable " + f"with CUTLASS Blackwell FNA Backward. " + "Try selecting a combination from: \n" + " natten.get_bwd_configs_for_blackwell_fna(q, k, v)" + "\n" + "Here's a few examples of available combinations for your use case:\n" + f"{examples}" + ) + + +def check_cutlass_blackwell_fmha_backward_config( + input_tensor: Tensor, + q_tile_size: Optional[int] = None, + kv_tile_size: Optional[int] = None, +) -> CutlassBlackwellFmhaBackwardConfigType: + assert input_tensor.dim() == 4 + + if (q_tile_size is None) ^ (kv_tile_size is None): + raise ValueError( + "Please specify both q_tile_size and kv_tile_size, or neither one. " + f"Got {q_tile_size=}, {kv_tile_size=}." + ) + + if q_tile_size is None and kv_tile_size is None: + q_tile_size, kv_tile_size = get_default_backward_tile_sizes( + input_tensor=input_tensor + ) + return (q_tile_size, kv_tile_size) + + tile_sizes = get_all_fmha_backward_configs(input_tensor=input_tensor) + + for q_t, kv_t in tile_sizes: + if q_t == q_tile_size and kv_t == kv_tile_size: + return (q_t, kv_t) + + # Fail and make suggestions + MAX_EXAMPLES = 3 + examples = "" + for i, (q_t, kv_t) in enumerate(tile_sizes): + examples += f"\n q_tile_size={q_t}, kv_tile_size={kv_t}" + if i > MAX_EXAMPLES: + break + + raise ValueError( + f"Invalid configuration for CUTLASS Blackwell FMHA Backward. " + f"Q tile size {q_tile_size} and KV tile size {kv_tile_size} " + f"are not among the {len(tile_sizes)} configurations implementable " + f"with CUTLASS Blackwell FMHA Backward. " + "Try selecting a combination from: \n" + " natten.get_bwd_configs_for_blackwell_fmha(q, k, v)" + "\n" + "Here's a few examples of available combinations for your use case:\n" + f"{examples}" + ) diff --git a/natten/torch-ext/natten/backends/configs/cutlass_hopper/__init__.py b/natten/torch-ext/natten/backends/configs/cutlass_hopper/__init__.py new file mode 100644 index 00000000..a22162df --- /dev/null +++ b/natten/torch-ext/natten/backends/configs/cutlass_hopper/__init__.py @@ -0,0 +1,522 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +from typing import List, Optional + +import torch # noqa: F401 +from torch import Tensor + +from ....types import ( + CutlassHopperFmhaBackwardConfigType, + CutlassHopperFmhaForwardConfigType, + CutlassHopperFnaBackwardConfigType, + CutlassHopperFnaForwardConfigType, + DimensionType, + KernelSchedule, +) +from ....utils.checks import check_tile_shape +from ....utils.device import get_device_cc + +DTYPE_TO_BITS = { + torch.float16: 16, + torch.bfloat16: 16, +} + +# TODO: notes + +HOPPER_FORWARD_CONFIGS = { + 1: { + 16: { + 32: [ + (((64,), (128,)), KernelSchedule.NonPersistent), + ], + 64: [ + (((64,), (128,)), KernelSchedule.NonPersistent), + ], + 128: [ + (((128,), (128,)), KernelSchedule.WarpSpecializedCooperative), + (((128,), (128,)), KernelSchedule.WarpSpecializedPingpong), + ], + 256: [ + (((128,), (64,)), KernelSchedule.WarpSpecializedCooperative), + ], + }, + }, + 2: { + 16: { + 32: [ + (((8, 8), (16, 8)), KernelSchedule.NonPersistent), + (((8, 8), (8, 16)), KernelSchedule.NonPersistent), + ], + 64: [ + (((8, 8), (16, 8)), KernelSchedule.NonPersistent), + (((8, 8), (8, 16)), KernelSchedule.NonPersistent), + ], + 128: [ + (((16, 8), (16, 8)), KernelSchedule.WarpSpecializedCooperative), + (((16, 8), (16, 8)), KernelSchedule.WarpSpecializedPingpong), + ], + 256: [ + (((16, 8), (8, 8)), KernelSchedule.WarpSpecializedCooperative), + (((8, 16), (8, 8)), KernelSchedule.WarpSpecializedCooperative), + ], + }, + }, + 3: { + 16: { + 32: [ + (((4, 4, 4), (4, 4, 8)), KernelSchedule.NonPersistent), + (((4, 4, 4), (2, 8, 8)), KernelSchedule.NonPersistent), + ], + 64: [ + (((4, 4, 4), (4, 4, 8)), KernelSchedule.NonPersistent), + (((4, 4, 4), (2, 8, 8)), KernelSchedule.NonPersistent), + ], + 128: [ + (((4, 4, 8), (4, 4, 8)), KernelSchedule.WarpSpecializedCooperative), + (((4, 4, 8), (4, 4, 8)), KernelSchedule.WarpSpecializedPingpong), + (((2, 8, 8), (2, 8, 8)), KernelSchedule.WarpSpecializedCooperative), + (((2, 8, 8), (2, 8, 8)), KernelSchedule.WarpSpecializedPingpong), + ], + 256: [ + (((4, 4, 8), (4, 4, 4)), KernelSchedule.WarpSpecializedCooperative), + (((2, 8, 8), (4, 4, 4)), KernelSchedule.WarpSpecializedCooperative), + ], + }, + }, +} + +HOPPER_BACKWARD_CONFIGS = { + 1: { + 16: { + 32: [ + ((64,), (128,)), + ((128,), (128,)), + ], + 64: [ + ((64,), (128,)), + ((128,), (128,)), + ], + 128: [ + ((64,), (128,)), + ], + }, + }, + 2: { + 16: { + 32: [ + ((8, 8), (16, 8)), + ((8, 8), (8, 16)), + ((16, 8), (16, 8)), + ((16, 8), (8, 16)), + ], + 64: [ + ((8, 8), (16, 8)), + ((8, 8), (8, 16)), + ((16, 8), (16, 8)), + ((16, 8), (8, 16)), + ], + 128: [ + ((8, 8), (16, 8)), + ((8, 8), (8, 16)), + ], + }, + }, + 3: { + 16: { + 32: [ + ((4, 4, 4), (4, 4, 8)), + ((4, 4, 4), (2, 8, 8)), + ((4, 4, 8), (4, 4, 8)), + ((4, 4, 8), (2, 8, 8)), + ], + 64: [ + ((4, 4, 4), (4, 4, 8)), + ((4, 4, 4), (2, 8, 8)), + ((4, 4, 8), (4, 4, 8)), + ((4, 4, 8), (2, 8, 8)), + ], + 128: [ + ((4, 4, 4), (4, 4, 8)), + ((4, 4, 4), (2, 8, 8)), + ((2, 4, 8), (2, 8, 8)), + ((1, 8, 8), (2, 8, 8)), + ], + }, + }, +} + + +def get_all_forward_configs( + input_tensor: Tensor, +) -> List[CutlassHopperFnaForwardConfigType]: + assert input_tensor.dim() in [4, 5, 6] + na_dim = input_tensor.dim() - 3 # batch, heads, head_dim + + dtype = input_tensor.dtype + dtype_bits = DTYPE_TO_BITS[dtype] + + head_dim = input_tensor.shape[-1] + + # if dtype not in [torch.float16, torch.bfloat16]: + if dtype_bits not in HOPPER_FORWARD_CONFIGS[na_dim]: # type: ignore + return [] + + # if head_dim not in [32, 64, 128, 256]: + if head_dim not in HOPPER_FORWARD_CONFIGS[na_dim][dtype_bits]: # type: ignore + return [] + + device_cc = get_device_cc(input_tensor.device) + if device_cc != 90: + return [] + + return HOPPER_FORWARD_CONFIGS[na_dim][dtype_bits][head_dim] # type: ignore + + +def get_all_backward_configs( + input_tensor: Tensor, +) -> List[CutlassHopperFnaBackwardConfigType]: + assert input_tensor.dim() in [4, 5, 6] + na_dim = input_tensor.dim() - 3 # batch, heads, head_dim + + dtype = input_tensor.dtype + dtype_bits = DTYPE_TO_BITS[dtype] + + head_dim = input_tensor.shape[-1] + + # if dtype not in [torch.float16, torch.bfloat16]: + if dtype_bits not in HOPPER_BACKWARD_CONFIGS[na_dim]: # type: ignore + return [] + + # if head_dim not in [32, 64, 128, 256]: + if head_dim not in HOPPER_BACKWARD_CONFIGS[na_dim][dtype_bits]: # type: ignore + return [] + + device_cc = get_device_cc(input_tensor.device) + if device_cc != 90: + return [] + + return HOPPER_BACKWARD_CONFIGS[na_dim][dtype_bits][head_dim] # type: ignore + + +# For FMHA +def get_all_fmha_forward_configs( + input_tensor: Tensor, +) -> List[CutlassHopperFmhaForwardConfigType]: + if input_tensor.dim() != 4: + raise ValueError("Only 4-D tensors are supported in FMHA.") + + configs_multi_dim = get_all_forward_configs(input_tensor) + assert all(len(q_t) == len(kv_t) == 1 for (q_t, kv_t), _ in configs_multi_dim) + + configs_fmha = [ + ((q_t[0], kv_t[0]), sched) for (q_t, kv_t), sched in configs_multi_dim + ] + + return configs_fmha + + +def get_all_fmha_backward_configs( + input_tensor: Tensor, +) -> List[CutlassHopperFmhaBackwardConfigType]: + if input_tensor.dim() != 4: + raise ValueError("Only 4-D tensors are supported in FMHA.") + + configs_multi_dim = get_all_backward_configs(input_tensor) + assert all(len(q_t) == len(kv_t) == 1 for q_t, kv_t in configs_multi_dim) + + configs_fmha = [(q_t[0], kv_t[0]) for q_t, kv_t in configs_multi_dim] + + return configs_fmha + + +def get_default_forward_config( + input_tensor: Tensor, +) -> CutlassHopperFnaForwardConfigType: + all_configs = get_all_forward_configs(input_tensor) + + if len(all_configs) < 1: + device_cc = get_device_cc(input_tensor.device) + raise ValueError( + "No configs exist for this use case; Hopper FMHA/FNA does not support it: " + f"{input_tensor.shape=}, {input_tensor.dtype=}, {device_cc=}." + ) + + return all_configs[0] + + +def get_default_backward_config( + input_tensor: Tensor, +) -> CutlassHopperFnaBackwardConfigType: + all_configs = get_all_backward_configs(input_tensor) + + if len(all_configs) < 1: + device_cc = get_device_cc(input_tensor.device) + raise ValueError( + "No configs exist for this use case; Hopper FMHA/FNA does not support it: " + f"{input_tensor.shape=}, {input_tensor.dtype=}, {device_cc=}." + ) + + return all_configs[0] + + +def get_default_fmha_forward_config( + input_tensor: Tensor, +) -> CutlassHopperFmhaForwardConfigType: + if input_tensor.dim() != 4: + raise ValueError("Only 4-D tensors are supported in FMHA.") + + (q_t, kv_t), sched = get_default_forward_config(input_tensor) + assert len(q_t) == len(kv_t) == 1 + + return (q_t[0], kv_t[0]), sched + + +def get_default_fmha_backward_config( + input_tensor: Tensor, +) -> CutlassHopperFmhaBackwardConfigType: + if input_tensor.dim() != 4: + raise ValueError("Only 4-D tensors are supported in FMHA.") + + q_t, kv_t = get_default_backward_config(input_tensor) + assert len(q_t) == len(kv_t) == 1 + + return q_t[0], kv_t[0] + + +def check_cutlass_hopper_fna_forward_config( + input_tensor: Tensor, + q_tile_shape: Optional[DimensionType] = None, + kv_tile_shape: Optional[DimensionType] = None, + kernel_schedule: Optional[KernelSchedule] = None, +) -> CutlassHopperFnaForwardConfigType: + assert input_tensor.dim() in [4, 5, 6] + na_dim = input_tensor.dim() - 3 # batch, heads, head_dim + + if (q_tile_shape is None) ^ (kv_tile_shape is None): + raise ValueError( + "Please specify both q_tile_shape and kv_tile_shape, or neither one. " + f"Got {q_tile_shape=}, {kv_tile_shape=}." + ) + + (default_q_tile_shape, default_kv_tile_shape), default_sched = ( + get_default_forward_config(input_tensor=input_tensor) + ) + if q_tile_shape is None and kv_tile_shape is None and kernel_schedule is None: + return (default_q_tile_shape, default_kv_tile_shape), default_sched # type: ignore[return-value] + + elif q_tile_shape is None and kv_tile_shape is None: + q_tile_shape = default_q_tile_shape + kv_tile_shape = default_kv_tile_shape + + q_tile_shape = check_tile_shape(q_tile_shape) + kv_tile_shape = check_tile_shape(kv_tile_shape) + + configs = get_all_forward_configs(input_tensor=input_tensor) + + for (q_t, kv_t), sched in configs: + if ( + q_t == q_tile_shape + and kv_t == kv_tile_shape + and (kernel_schedule is None or sched == kernel_schedule) + ): + return (q_t, kv_t), sched # type: ignore + + # Fail and make suggestions + MAX_EXAMPLES = 3 + examples = "" + for i, ((q_t, kv_t), sched) in enumerate(configs): + examples += f"\n q_tile_shape={q_t}, kv_tile_shape={kv_t}, schedule={sched}" + if i > MAX_EXAMPLES: + break + + raise ValueError( + f"Invalid configuration for CUTLASS Hopper FNA-{na_dim}D. " + f"Q tile shape {q_tile_shape}, KV tile shape {kv_tile_shape}, and schedule {kernel_schedule} " + f"are not among the {len(configs)} configurations implementable " + f"with CUTLASS Hopper FNA. " + "Try selecting a combination from: \n" + " natten.get_configs_for_hopper_fna(q, k, v)" + "\n" + "Here's a few examples of available combinations for your use case:\n" + f"{examples}" + ) + + +def check_cutlass_hopper_fna_backward_config( + input_tensor: Tensor, + q_tile_shape: Optional[DimensionType] = None, + kv_tile_shape: Optional[DimensionType] = None, +) -> CutlassHopperFnaBackwardConfigType: + assert input_tensor.dim() in [4, 5, 6] + na_dim = input_tensor.dim() - 3 # batch, heads, head_dim + + if (q_tile_shape is None) ^ (kv_tile_shape is None): + raise ValueError( + "Please specify both q_tile_shape and kv_tile_shape, or neither one. " + f"Got {q_tile_shape=}, {kv_tile_shape=}." + ) + + default_q_tile_shape, default_kv_tile_shape = get_default_backward_config( + input_tensor=input_tensor + ) + if q_tile_shape is None and kv_tile_shape is None: + return default_q_tile_shape, default_kv_tile_shape # type: ignore[return-value] + + elif q_tile_shape is None and kv_tile_shape is None: + q_tile_shape = default_q_tile_shape + kv_tile_shape = default_kv_tile_shape + + q_tile_shape = check_tile_shape(q_tile_shape) + kv_tile_shape = check_tile_shape(kv_tile_shape) + + configs = get_all_backward_configs(input_tensor=input_tensor) + + for q_t, kv_t in configs: + if q_t == q_tile_shape and kv_t == kv_tile_shape: + return q_t, kv_t # type: ignore + + # Fail and make suggestions + MAX_EXAMPLES = 3 + examples = "" + for i, (q_t, kv_t) in enumerate(configs): + examples += f"\n q_tile_shape={q_t}, kv_tile_shape={kv_t}" + if i > MAX_EXAMPLES: + break + + raise ValueError( + f"Invalid configuration for CUTLASS Hopper FNA-{na_dim}D Backward. " + f"Q tile shape {q_tile_shape} and KV tile shape {kv_tile_shape} " + f"are not among the {len(configs)} configurations implementable " + f"with CUTLASS Hopper FNA Backward. " + "Try selecting a combination from: \n" + " natten.get_bwd_configs_for_hopper_fna(q, k, v)" + "\n" + "Here's a few examples of available combinations for your use case:\n" + f"{examples}" + ) + + +def check_cutlass_hopper_fmha_forward_config( + input_tensor: Tensor, + q_tile_size: Optional[int] = None, + kv_tile_size: Optional[int] = None, + kernel_schedule: Optional[KernelSchedule] = None, +) -> CutlassHopperFmhaForwardConfigType: + assert input_tensor.dim() == 4 + + if (q_tile_size is None) ^ (kv_tile_size is None): + raise ValueError( + "Please specify both q_tile_size and kv_tile_size, or neither one. " + f"Got {q_tile_size=}, {kv_tile_size=}." + ) + + (default_q_tile_size, default_kv_tile_size), default_sched = ( + get_default_fmha_forward_config(input_tensor=input_tensor) + ) + if q_tile_size is None and kv_tile_size is None and kernel_schedule is None: + return (default_q_tile_size, default_kv_tile_size), default_sched + + elif q_tile_size is None and kv_tile_size is None: + q_tile_size = default_q_tile_size + kv_tile_size = default_kv_tile_size + + configs = get_all_fmha_forward_configs(input_tensor=input_tensor) + + for (q_t, kv_t), sched in configs: + if ( + q_t == q_tile_size + and kv_t == kv_tile_size + and (kernel_schedule is None or sched == kernel_schedule) + ): + return (q_t, kv_t), sched # type: ignore + + # Fail and make suggestions + MAX_EXAMPLES = 3 + examples = "" + for i, ((q_t, kv_t), sched) in enumerate(configs): + examples += f"\n q_tile_size={q_t}, kv_tile_size={kv_t}, schedule={sched}" + if i > MAX_EXAMPLES: + break + + raise ValueError( + f"Invalid configuration for CUTLASS Hopper FMHA. " + f"Q tile size {q_tile_size}, KV tile size {kv_tile_size}, and schedule {kernel_schedule} " + f"are not among the {len(configs)} configurations implementable " + f"with CUTLASS Hopper FMHA. " + "Try selecting a combination from: \n" + " natten.get_configs_for_hopper_fmha(q, k, v)" + "\n" + "Here's a few examples of available combinations for your use case:\n" + f"{examples}" + ) + + +def check_cutlass_hopper_fmha_backward_config( + input_tensor: Tensor, + q_tile_size: Optional[int] = None, + kv_tile_size: Optional[int] = None, +) -> CutlassHopperFmhaBackwardConfigType: + assert input_tensor.dim() == 4 + + if (q_tile_size is None) ^ (kv_tile_size is None): + raise ValueError( + "Please specify both q_tile_size and kv_tile_size, or neither one. " + f"Got {q_tile_size=}, {kv_tile_size=}." + ) + + default_q_tile_size, default_kv_tile_size = get_default_fmha_backward_config( + input_tensor=input_tensor + ) + if q_tile_size is None and kv_tile_size is None: + return default_q_tile_size, default_kv_tile_size + + elif q_tile_size is None and kv_tile_size is None: + q_tile_size = default_q_tile_size + kv_tile_size = default_kv_tile_size + + configs = get_all_fmha_backward_configs(input_tensor=input_tensor) + + for q_t, kv_t in configs: + if q_t == q_tile_size and kv_t == kv_tile_size: + return q_t, kv_t # type: ignore + + # Fail and make suggestions + MAX_EXAMPLES = 3 + examples = "" + for i, (q_t, kv_t) in enumerate(configs): + examples += f"\n q_tile_size={q_t}, kv_tile_size={kv_t}" + if i > MAX_EXAMPLES: + break + + raise ValueError( + f"Invalid configuration for CUTLASS Hopper FMHA Backward. " + f"Q tile size {q_tile_size} and KV tile size {kv_tile_size} " + f"are not among the {len(configs)} configurations implementable " + f"with CUTLASS Hopper FMHA Backward. " + "Try selecting a combination from: \n" + " natten.get_bwd_configs_for_hopper_fmha(q, k, v)" + "\n" + "Here's a few examples of available combinations for your use case:\n" + f"{examples}" + ) diff --git a/natten/torch-ext/natten/backends/configs/flex/__init__.py b/natten/torch-ext/natten/backends/configs/flex/__init__.py new file mode 100644 index 00000000..0c0e24ee --- /dev/null +++ b/natten/torch-ext/natten/backends/configs/flex/__init__.py @@ -0,0 +1,210 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +from typing import List, Optional + +import torch # noqa: F401 +from torch import Tensor + +from ....types import ( + DimensionType, + FlexFmhaForwardConfigType, + FlexFnaForwardConfigType, +) +from ....utils.checks import check_tile_shape +from ....utils.device import get_device_cc + +# TODO: add more tile sizes/shapes +# TODO: add backprop tile sizes/shapes +# Only doing 64 x 64 for now, since it's the one that successfully compiles across devices and +# use cases without running into compile errors (i.e. shmem over-subscription) +# Once Flex with compilation actually starts working as expected and is out of prototype, we can +# add in more tile sizes/shapes and condition them on arch / use case, like we do for CUTLASS FNA. + +FLEX_FORWARD_TILE_SHAPES = { + 1: [ + # ((128, ), (128, )), + ((64,), (64,)), + ], + 2: [ + # ((8, 16), (8, 16)), + ((8, 8), (8, 8)), + ((4, 16), (4, 16)), + ((4, 16), (8, 8)), + ], + 3: [ + # ((4, 4, 8), (4, 4, 8)), + ((4, 4, 4), (4, 4, 4)), + ((2, 4, 8), (2, 4, 8)), + ((2, 4, 8), (4, 4, 4)), + ], +} + + +def _get_default_tile_shapes_forward( + na_dim: int, +) -> FlexFnaForwardConfigType: + assert na_dim in [1, 2, 3] + + if na_dim == 1: + return ((64,), (64,)) + if na_dim == 2: + return ((8, 8), (8, 8)) + if na_dim == 3: + return ((4, 4, 4), (4, 4, 4)) + + raise NotImplementedError() + + +def get_all_tile_shapes_forward( + input_tensor: Tensor, +) -> List[FlexFnaForwardConfigType]: + assert input_tensor.dim() in [4, 5, 6] + na_dim = input_tensor.dim() - 3 # batch, heads, head_dim + + return FLEX_FORWARD_TILE_SHAPES[na_dim] # type: ignore + + +# For FMHA +def get_all_tile_sizes_forward(input_tensor: Tensor) -> List[FlexFmhaForwardConfigType]: + if input_tensor.dim() != 4: + raise ValueError("Only 4-D tensors are supported in FMHA.") + + tile_shapes = get_all_tile_shapes_forward(input_tensor) + assert all(len(q_t) == len(kv_t) == 1 for q_t, kv_t in tile_shapes) + + tile_sizes = [(q_t[0], kv_t[0]) for q_t, kv_t in tile_shapes] + + return tile_sizes + + +def get_default_forward_tile_shapes(input_tensor: Tensor) -> FlexFnaForwardConfigType: + assert input_tensor.dim() in [4, 5, 6] + na_dim = input_tensor.dim() - 3 # batch, heads, head_dim + + return _get_default_tile_shapes_forward(na_dim) + + +def get_default_forward_tile_sizes(input_tensor: Tensor) -> FlexFmhaForwardConfigType: + if input_tensor.dim() != 4: + raise ValueError("Only 4-D tensors are supported in FMHA.") + + q_t, kv_t = get_default_forward_tile_shapes(input_tensor) + assert len(q_t) == len(kv_t) == 1 + + return (q_t[0], kv_t[0]) + + +def check_flex_fna_forward_config( + input_tensor: Tensor, + q_tile_shape: Optional[DimensionType] = None, + kv_tile_shape: Optional[DimensionType] = None, +) -> FlexFnaForwardConfigType: + assert input_tensor.dim() in [4, 5, 6] + na_dim = input_tensor.dim() - 3 # batch, heads, head_dim + + if (q_tile_shape is None) ^ (kv_tile_shape is None): + raise ValueError( + "Please specify both q_tile_shape and kv_tile_shape, or neither one. " + f"Got {q_tile_shape=}, {kv_tile_shape=}." + ) + + if q_tile_shape is None and kv_tile_shape is None: + return get_default_forward_tile_shapes(input_tensor=input_tensor) + + q_tile_shape = check_tile_shape(q_tile_shape) + kv_tile_shape = check_tile_shape(kv_tile_shape) + + tile_shapes = get_all_tile_shapes_forward(input_tensor=input_tensor) + + for q_t, kv_t in tile_shapes: + if q_t == q_tile_shape and kv_t == kv_tile_shape: + return (q_t, kv_t) # type: ignore + + # Fail and make suggestions + device_cc = get_device_cc(input_tensor.device) + MAX_EXAMPLES = 3 + examples = "" + for i, (q_t, kv_t) in enumerate(tile_shapes): + examples += f"\n q_tile_shape={q_t}, kv_tile_shape={kv_t}" + if i > MAX_EXAMPLES: + break + + raise ValueError( + f"Invalid configuration for Flex FNA-{na_dim}D. " + f"Q tile shape {q_tile_shape} and KV tile shape {kv_tile_shape} " + f"are not among the {len(tile_shapes)} configurations implementable " + f"with Flex FNA for SM{device_cc}. " + "Try selecting a combination from: \n" + " natten.get_configs_for_flex_fna(q, k, v)" + "\n" + "Here's a few examples of available combinations for your use case:\n" + f"{examples}" + ) + + +def check_flex_fmha_forward_config( + input_tensor: Tensor, + q_tile_size: Optional[int] = None, + kv_tile_size: Optional[int] = None, +) -> FlexFmhaForwardConfigType: + assert input_tensor.dim() == 4 + + if (q_tile_size is None) ^ (kv_tile_size is None): + raise ValueError( + "Please specify both q_tile_size and kv_tile_size, or neither one. " + f"Got {q_tile_size=}, {kv_tile_size=}." + ) + + if q_tile_size is None and kv_tile_size is None: + q_tile_size, kv_tile_size = get_default_forward_tile_sizes( + input_tensor=input_tensor + ) + return (q_tile_size, kv_tile_size) + + tile_sizes = get_all_tile_sizes_forward(input_tensor=input_tensor) + + for q_t, kv_t in tile_sizes: + if q_t == q_tile_size and kv_t == kv_tile_size: + return (q_t, kv_t) + + # Fail and make suggestions + device_cc = get_device_cc(input_tensor.device) + MAX_EXAMPLES = 3 + examples = "" + for i, (q_t, kv_t) in enumerate(tile_sizes): + examples += f"\n q_tile_size={q_t}, kv_tile_size={kv_t}" + if i > MAX_EXAMPLES: + break + + raise ValueError( + f"Invalid configuration for Flex FMHA. " + f"Q tile size {q_tile_size} and KV tile size {kv_tile_size} " + f"are not among the {len(tile_sizes)} configurations implementable " + f"with Flex FMHA for SM{device_cc}. " + "Try selecting a combination from: \n" + " natten.get_configs_for_flex_fmha(q, k, v)\n" + "\n" + "Here's a few examples of available combinations for your use case:\n" + f"{examples}" + ) diff --git a/natten/torch-ext/natten/backends/flex.py b/natten/torch-ext/natten/backends/flex.py new file mode 100644 index 00000000..5ee3b071 --- /dev/null +++ b/natten/torch-ext/natten/backends/flex.py @@ -0,0 +1,799 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +import functools +import math +import time +from typing import Callable, Optional, Tuple, Union + +import torch +from torch import BoolTensor, IntTensor, Tensor +from torch.nn.attention.flex_attention import ( + BlockMask, + create_block_mask, + flex_attention, +) + +from ..backends.configs.checks import ( # noqa: F401 + _FLEX_COMPILE_SUPPORTED, + _FLEX_SUPPORTED, + can_run_flex_attention, +) +from ..backends.configs.flex import ( + check_flex_fmha_forward_config, + check_flex_fna_forward_config, +) +from ..token_permute import token_permute_operation, token_unpermute_operation +from ..types import ( + CausalArg1DTypeOrDed, + CausalArg2DTypeOrDed, + CausalArg3DTypeOrDed, + CausalArgType, + CausalArgTypeOrDed, + Dimension1DType, + Dimension1DTypeOrDed, + Dimension2DType, + Dimension2DTypeOrDed, + Dimension3DType, + Dimension3DTypeOrDed, + DimensionType, + DimensionTypeOrDed, +) +from ..utils import log +from ..utils.checks import ( + check_all_args, + check_args_against_input, + check_input_size_arg, + fmha_tensor_checks, + na_tensor_checks, + varlen_tensor_checks, +) +from ..utils.environment import is_torch_compiling + +logger = log.get_logger(__name__) + + +def get_flex_attention_fn( + torch_compile: bool, torch_compile_args: Optional[dict] = None +) -> Callable: + if not torch_compile: + return flex_attention + + additional_args = torch_compile_args or {} + additional_args["dynamic"] = False + + return torch.compile(flex_attention, **additional_args) + + +def _run_flex_attn( + q: Tensor, + k: Tensor, + v: Tensor, + block_mask: BlockMask, + scale: float, + torch_compile: bool, + q_tile_size: Optional[int] = None, + kv_tile_size: Optional[int] = None, + torch_compile_args: Optional[dict] = None, +) -> Tuple[Tensor, Tensor]: + + # We may need to override the default flex config. + # Default ones are not guaranteed to work out of the box across architectures. + # Some oversubscribe shmem even on the B200! + torch_compile_args = {} + + # Disable flex decoding path + kernel_options = { + "FORCE_USE_FLEX_ATTENTION": True, + } + + if q_tile_size is not None and torch_compile: + kv_tile_size = kv_tile_size or q_tile_size + + # Have to auto-tune, otherwise torch will only allow the default config. + torch_compile_args["mode"] = "max-autotune-no-cudagraphs" + + kernel_options["SPARSE_Q_BLOCK_SIZE"] = q_tile_size # type: ignore[assignment] + kernel_options["SPARSE_KV_BLOCK_SIZE"] = kv_tile_size # type: ignore[assignment] + kernel_options["BLOCK_M"] = q_tile_size # type: ignore[assignment] + kernel_options["BLOCK_N"] = kv_tile_size # type: ignore[assignment] + + flex_fn = get_flex_attention_fn( + torch_compile=torch_compile, torch_compile_args=torch_compile_args + ) + + # tensors are BHSD here + is_gqa = q.shape[1] != k.shape[1] + return flex_fn( + q, + k, + v, + block_mask=block_mask, + return_lse=True, + scale=scale, + kernel_options=kernel_options, + enable_gqa=is_gqa, + ) + + +def run_flex_attn( + q: Tensor, + k: Tensor, + v: Tensor, + block_mask: BlockMask, + scale: float, + torch_compile: bool, + torch_compile_args: Optional[dict] = None, + q_tile_size: Optional[int] = None, + kv_tile_size: Optional[int] = None, +) -> Tuple[Tensor, Tensor]: + + if q_tile_size is not None and kv_tile_size is not None: + return _run_flex_attn( + q, + k, + v, + block_mask=block_mask, + scale=scale, + torch_compile=torch_compile, + q_tile_size=q_tile_size, + kv_tile_size=kv_tile_size, + torch_compile_args=torch_compile_args, + ) + + # Use smallest tile size combo to try and evade shmem oversubscription + # The defaults just fail very frequently. + return _run_flex_attn( + q, + k, + v, + block_mask=block_mask, + scale=scale, + torch_compile=torch_compile, + q_tile_size=64, + kv_tile_size=64, + torch_compile_args=torch_compile_args, + ) + + +def flex_fmha( + query: Tensor, + key: Tensor, + value: Tensor, + is_causal: bool = False, + scale: Optional[float] = None, + q_tile_size: Optional[int] = None, + kv_tile_size: Optional[int] = None, + torch_compile: bool = False, + return_lse: bool = False, + # varlen parameters + cumulative_seqlen_Q: Optional[Tensor] = None, + cumulative_seqlen_KV: Optional[Tensor] = None, + max_seqlen_Q: int = 0, + max_seqlen_KV: int = 0, +) -> Union[Tensor, Tuple[Tensor, Tensor]]: + + fmha_tensor_checks( + query, + key, + value, + must_match_head_dims=True, + supports_gqa_mqa=True, + backend_name="Flex FMHA", + ) + + ( + cumulative_seqlen_Q, + cumulative_seqlen_KV, + max_seqlen_Q, + max_seqlen_KV, + ) = varlen_tensor_checks( + query=query, + key=key, + value=value, + cumulative_seqlen_Q=cumulative_seqlen_Q, + cumulative_seqlen_KV=cumulative_seqlen_KV, + max_seqlen_Q=max_seqlen_Q, + max_seqlen_KV=max_seqlen_KV, + ) + is_varlen = cumulative_seqlen_Q is not None + + assert can_run_flex_attention( + query, + key, + value, + is_causal=is_causal, + is_varlen=is_varlen, + torch_compile=torch_compile, + raise_error=True, + ) + + q_tile_size, kv_tile_size = check_flex_fmha_forward_config( + input_tensor=query, + q_tile_size=q_tile_size, + kv_tile_size=kv_tile_size, + ) + + scale = scale or query.shape[-1] ** -0.5 + + batch_size, seqlen_q, num_heads, head_dim = query.shape + _, seqlen_kv, num_heads_kv, head_dim_v = value.shape + + # Flex and torch attention use heads first layout + query_ = query.reshape(batch_size, seqlen_q, num_heads, head_dim).transpose(1, 2) + key_ = key.reshape(batch_size, seqlen_kv, num_heads_kv, head_dim_v).transpose(1, 2) + value_ = value.reshape(batch_size, seqlen_kv, num_heads_kv, head_dim_v).transpose( + 1, 2 + ) + + out_, lse_ = run_flex_attn( + query_, + key_, + value_, + block_mask=None, # type: ignore[arg-type] + scale=scale, + torch_compile=torch_compile, + q_tile_size=q_tile_size, + kv_tile_size=kv_tile_size, + ) + + out = out_.transpose(1, 2).reshape(batch_size, seqlen_q, num_heads, head_dim_v) + lse = lse_.transpose(1, 2).reshape(batch_size, seqlen_q, num_heads) + + if return_lse: + return out, lse + + return out + + +# TODO: move me elsewhere? +def idx2crd(index, shape) -> tuple: + rank = len(shape) + coord = [] + residual = index + for i in range(rank - 1, -1, -1): + coord.append(residual % shape[i]) + residual = residual // shape[i] + + # assert residual == 0 + return tuple(coord[::-1]) + + +def get_na_flex_mask( + device: str, + na_dim: int, + qkv_shape: DimensionType, + kernel_size: DimensionType, + stride: DimensionType, + dilation: DimensionType, + is_causal: CausalArgType, + q_tile_shape: Optional[DimensionType] = None, + kv_tile_shape: Optional[DimensionType] = None, + q_shape: Optional[DimensionType] = None, + kv_shape: Optional[DimensionType] = None, + torch_compile: bool = False, +): + num_dilation_groups = math.prod(dilation) + if not is_torch_compiling(): + flex_mask_start_time = time.perf_counter() + do_token_permute = q_tile_shape is not None and kv_tile_shape is not None + if do_token_permute: + if q_tile_shape is None or kv_tile_shape is None: + raise ValueError( + "Please specify Q and KV tile shapes for multi dimensional tiling. " + f"Got {q_tile_shape=}, {kv_tile_shape=}." + ) + + if q_shape is None or kv_shape is None: + raise ValueError( + "Please specify q_shape and kv_shape for multi dimensional tiling." + ) + + if len(q_tile_shape) != na_dim or len(kv_tile_shape) != na_dim: + raise ValueError( + "Q and KV tile shapes must match the number of dimensions in the " + f"token layout ({na_dim}, got {q_tile_shape=}, {kv_tile_shape=}." + ) + + if any(x % t != 0 for x, t in zip(q_shape, q_tile_shape)): + raise ValueError( + "Input must be divisible by Q tile shape, but got " + f"{q_shape=}, {q_tile_shape=}." + ) + + if any(x % t != 0 for x, t in zip(kv_shape, kv_tile_shape)): + raise ValueError( + "Input must be divisible by KV tile shape, but got " + f"{kv_shape=}, {kv_tile_shape=}." + ) + + q_rest_shape = tuple(x // t for x, t in zip(q_shape, q_tile_shape)) + kv_rest_shape = tuple(x // t for x, t in zip(kv_shape, kv_tile_shape)) + + def single_dim_tiling_mask( + b: IntTensor, + h: IntTensor, + q_idx: IntTensor, + kv_idx: IntTensor, + qkv_shape, + kernel_size, + stride, + dilation, + is_causal, + ) -> BoolTensor: + + # Reconstruct global Q and KV coordinates + q_crd = idx2crd(q_idx, qkv_shape) + kv_crd = idx2crd(kv_idx, qkv_shape) + + # Coordinates within dilation group + q_crd_di = tuple(x // d for x, d in zip(q_crd, dilation)) + kv_crd_di = tuple(x // d for x, d in zip(kv_crd, dilation)) + + # Dilation group coordinates + q_dilation_group_crd = tuple(x % d for x, d in zip(q_crd, dilation)) + kv_dilation_group_crd = tuple(x % d for x, d in zip(kv_crd, dilation)) + + # Fixup input shape according to dilation group + dilation_group_padding = tuple( + 1 - ((dg + (d - (x % d))) // d) + for dg, d, x in zip(q_dilation_group_crd, dilation, qkv_shape) + ) + qkv_shape_corrected = tuple( + (x // d) + p for p, d, x in zip(dilation_group_padding, dilation, qkv_shape) + ) + + # Window size left and right (non-causal only) + window_size_left = tuple(w // 2 for w in kernel_size) + window_size_right = tuple(w // 2 + (w % 2 - 1) for w in kernel_size) + + masks = [] + for i in range(na_dim): + if is_causal[i]: + # Leader is the last (right-most) query in the stride group. + stride_group_leader = torch.min( + (q_crd_di[i] // stride[i]) * stride[i] + stride[i] - 1, + qkv_shape_corrected[i] - 1, + ) + + mask = ( + ( + q_crd_di[i] - kv_crd_di[i] >= 0 + ) # window still ends at query index + & (stride_group_leader - kv_crd_di[i] < kernel_size[i]) + & (q_dilation_group_crd[i] == kv_dilation_group_crd[i]) + ) + else: + # Leader is the center-most query in the stride group. + # If stride is even, choose the right hand side center query. + stride_group_leader = torch.min( + (q_crd_di[i] // stride[i]) * stride[i] + (stride[i] // 2), + qkv_shape_corrected[i] - 1, + ) + + window_center = stride_group_leader.clamp( + window_size_left[i] * torch.ones_like(qkv_shape_corrected[i]), + qkv_shape_corrected[i] - 1 - window_size_right[i], + ) + w0 = window_center - kv_crd_di[i] + w1 = kv_crd_di[i] - window_center + mask = ( + ((0 <= w0) & (w0 <= window_size_left[i])) + | ((0 <= w1) & (w1 <= window_size_right[i])) + ) & (q_dilation_group_crd[i] == kv_dilation_group_crd[i]) + + masks.append(mask) + + return functools.reduce(lambda x, y: x & y, masks) # type: ignore + + def multi_dim_tiling_mask( + b: IntTensor, + h: IntTensor, + q_idx: IntTensor, + kv_idx: IntTensor, + q_tile_size: int, + kv_tile_size: int, + q_tile_shape, + kv_tile_shape, + qkv_shape, + kernel_size, + stride, + dilation, + is_causal, + ) -> BoolTensor: + + # Reconstruct global Q and KV coordinates + q_tile_idx = q_idx // q_tile_size + kv_tile_idx = kv_idx // kv_tile_size + q_tile_offset = q_idx % q_tile_size + kv_tile_offset = kv_idx % q_tile_size + q_tile_coord = idx2crd(q_tile_idx, q_rest_shape) + kv_tile_coord = idx2crd(kv_tile_idx, kv_rest_shape) + q_tile_offset_coord = idx2crd(q_tile_offset, q_tile_shape) + kv_tile_offset_coord = idx2crd(kv_tile_offset, kv_tile_shape) + + q_crd = tuple( + tile_crd * tile_sz + tile_off + for tile_crd, tile_sz, tile_off in zip( + q_tile_coord, q_tile_shape, q_tile_offset_coord + ) + ) + kv_crd = tuple( + tile_crd * tile_sz + tile_off + for tile_crd, tile_sz, tile_off in zip( + kv_tile_coord, kv_tile_shape, kv_tile_offset_coord + ) + ) + + # Dilation group coordinates + # b_actual = b // num_dilation_groups + dilation_group_idx = b % num_dilation_groups + dilation_group_crd = idx2crd(dilation_group_idx, dilation) + + # Fixup input shape according to dilation group + dilation_group_padding = tuple( + 1 - ((dg + (d - (x % d))) // d) + for dg, d, x in zip(dilation_group_crd, dilation, qkv_shape) + ) + qkv_shape_corrected = tuple( + (x // d) + p for p, d, x in zip(dilation_group_padding, dilation, qkv_shape) + ) + + # Window size left and right (non-causal only) + window_size_left = tuple(w // 2 for w in kernel_size) + window_size_right = tuple(w // 2 + (w % 2 - 1) for w in kernel_size) + + masks = [] + for i in range(na_dim): + if is_causal[i]: + # Leader is the last (right-most) query in the stride group. + stride_group_leader = torch.min( + (q_crd[i] // stride[i]) * stride[i] + stride[i] - 1, + qkv_shape_corrected[i] - 1, + ) + + mask = ( + q_crd[i] - kv_crd[i] >= 0 + ) & ( # window still ends at query index + stride_group_leader - kv_crd[i] < kernel_size[i] + ) + else: + # Leader is the center-most query in the stride group. + # If stride is even, choose the right hand side center query. + stride_group_leader = torch.min( + (q_crd[i] // stride[i]) * stride[i] + (stride[i] // 2), + qkv_shape_corrected[i] - 1, + ) + + window_center = stride_group_leader.clamp( + window_size_left[i] * torch.ones_like(qkv_shape_corrected[i]), + qkv_shape_corrected[i] - 1 - window_size_right[i], + ) + w0 = window_center - kv_crd[i] + w1 = kv_crd[i] - window_center + + mask = ((0 <= w0) & (w0 <= window_size_left[i])) | ( + (0 <= w1) & (w1 <= window_size_right[i]) + ) + + masks.append(mask) + + return functools.reduce(lambda x, y: x & y, masks) # type: ignore + + mask_mod = None + seq_length_q = seq_length_kv = math.prod(qkv_shape) + q_tile_size, kv_tile_size = 64, 64 + if do_token_permute: + assert q_shape is not None + assert kv_shape is not None + assert q_tile_shape is not None + assert kv_tile_shape is not None + + seq_length_q = math.prod(q_shape) + seq_length_kv = math.prod(kv_shape) + q_tile_size, kv_tile_size = math.prod(q_tile_shape), math.prod(kv_tile_shape) + + mask_mod = functools.partial( + multi_dim_tiling_mask, + q_tile_size=q_tile_size, + kv_tile_size=kv_tile_size, + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + qkv_shape=qkv_shape, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + ) + else: + mask_mod = functools.partial( + single_dim_tiling_mask, + qkv_shape=qkv_shape, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + ) + + block_mask = create_block_mask( + mask_mod, + B=None, + H=None, + Q_LEN=seq_length_q, + KV_LEN=seq_length_kv, + _compile=torch_compile, + BLOCK_SIZE=(q_tile_size, kv_tile_size), + device=device, + ) + if not is_torch_compiling(): + flex_mask_end_time = time.perf_counter() + flex_mask_time = flex_mask_end_time - flex_mask_start_time + logger.debug( + f"Flex Attention block mask ({torch_compile=}) created in {flex_mask_time:.2f} seconds." + ) + return block_mask + + +def flex_fna_generic( + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size: DimensionTypeOrDed, + stride: DimensionTypeOrDed = 1, + dilation: DimensionTypeOrDed = 1, + is_causal: Optional[CausalArgTypeOrDed] = False, + scale: Optional[float] = None, + q_tile_shape: Optional[DimensionType] = None, + kv_tile_shape: Optional[DimensionType] = None, + torch_compile: bool = False, + return_lse: bool = False, +) -> Union[Tensor, Tuple[Tensor, Tensor]]: + + na_tensor_checks( + query, key, value, must_match_head_dims=True, supports_gqa_mqa=True + ) + + na_dim = query.dim() - 3 # batch, heads, head_dim + + kernel_size, stride, dilation, is_causal = check_all_args( + na_dim, kernel_size, stride, dilation, is_causal + ) + + check_args_against_input( + query, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + ) + + batch_size, *qkv_shape_in, num_heads, head_dim = query.shape + num_heads_kv, head_dim_v = value.shape[-2:] + qkv_shape = check_input_size_arg(na_dim, qkv_shape_in) + + scale = scale or query.shape[-1] ** -0.5 + + assert can_run_flex_attention( + query, key, value, torch_compile=torch_compile, raise_error=True + ) + + if (q_tile_shape is None) ^ (kv_tile_shape is None): + raise ValueError( + "Please specify both q_tile_shape and kv_tile_shape, or neither one. " + f"Got {q_tile_shape=}, {kv_tile_shape=}." + ) + + do_token_permute = q_tile_shape is not None and kv_tile_shape is not None + + q_shape = kv_shape = qkv_shape + q_tile_size: Optional[int] = None + kv_tile_size: Optional[int] = None + if do_token_permute: + q_tile_shape, kv_tile_shape = check_flex_fna_forward_config( + input_tensor=query, + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + ) + + q_tile_size = math.prod(q_tile_shape) + kv_tile_size = math.prod(kv_tile_shape) + + assert q_tile_shape is not None + assert kv_tile_shape is not None + query_perm, _, q_shape = token_permute_operation( + query, + tile_shape=q_tile_shape, + dilation=dilation, + flip_tiled_dims=False, + ) + key_perm, _, k_shape = token_permute_operation( + key, tile_shape=kv_tile_shape, dilation=dilation, flip_tiled_dims=False + ) + value_perm, _, v_shape = token_permute_operation( + value, tile_shape=kv_tile_shape, dilation=dilation, flip_tiled_dims=False + ) + + assert k_shape == v_shape + kv_shape = k_shape + + # Token permute already flattens to 1-D + # Flex uses heads first layout + query_ = query_perm.transpose(1, 2) + key_ = key_perm.transpose(1, 2) + value_ = value_perm.transpose(1, 2) + + else: + seqlen = math.prod(qkv_shape) + # Flex uses heads first layout + query_ = query.reshape(batch_size, seqlen, num_heads, head_dim).transpose(1, 2) + key_ = key.reshape(batch_size, seqlen, num_heads_kv, head_dim).transpose(1, 2) + value_ = value.reshape(batch_size, seqlen, num_heads_kv, head_dim_v).transpose( + 1, 2 + ) + + na_block_mask = get_na_flex_mask( + device=query.device.type, + na_dim=na_dim, + qkv_shape=qkv_shape, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + q_shape=q_shape, + kv_shape=kv_shape, + torch_compile=torch_compile, + ) + + out_, lse_ = run_flex_attn( + query_, + key_, + value_, + na_block_mask, + scale, + torch_compile=torch_compile, + q_tile_size=q_tile_size, + kv_tile_size=kv_tile_size, + ) + + if do_token_permute: + out = out_.transpose(1, 2) + lse = lse_.transpose(1, 2).unsqueeze(-1) + + assert q_tile_shape is not None + assert kv_tile_shape is not None + out = token_unpermute_operation( + out, + token_layout_shape=qkv_shape, + tile_shape=q_tile_shape, + dilation=dilation, + flip_tiled_dims=False, + ) + lse = token_unpermute_operation( + lse, + token_layout_shape=qkv_shape, + tile_shape=q_tile_shape, + dilation=dilation, + flip_tiled_dims=False, + ).squeeze(-1) + else: + out = out_.transpose(1, 2).reshape( + batch_size, *qkv_shape, num_heads, head_dim_v + ) + lse = lse_.transpose(1, 2).reshape(batch_size, *qkv_shape, num_heads) + + if return_lse: + return out, lse + + return out + + +def na1d_flex( + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size: Dimension1DTypeOrDed, + stride: Dimension1DTypeOrDed = 1, + dilation: Dimension1DTypeOrDed = 1, + is_causal: Optional[CausalArg1DTypeOrDed] = False, + scale: Optional[float] = None, + q_tile_shape: Optional[Dimension1DType] = None, + kv_tile_shape: Optional[Dimension1DType] = None, + torch_compile: bool = False, + return_lse: bool = False, +) -> Union[Tensor, Tuple[Tensor, Tensor]]: + return flex_fna_generic( + query=query, + key=key, + value=value, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + scale=scale, + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + torch_compile=torch_compile, + return_lse=return_lse, + ) + + +def na2d_flex( + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size: Dimension2DTypeOrDed, + stride: Dimension2DTypeOrDed = 1, + dilation: Dimension2DTypeOrDed = 1, + is_causal: Optional[CausalArg2DTypeOrDed] = False, + scale: Optional[float] = None, + q_tile_shape: Optional[Dimension2DType] = None, + kv_tile_shape: Optional[Dimension2DType] = None, + torch_compile: bool = False, + return_lse: bool = False, +) -> Union[Tensor, Tuple[Tensor, Tensor]]: + return flex_fna_generic( + query=query, + key=key, + value=value, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + scale=scale, + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + torch_compile=torch_compile, + return_lse=return_lse, + ) + + +def na3d_flex( + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size: Dimension3DTypeOrDed, + stride: Dimension3DTypeOrDed = 1, + dilation: Dimension3DTypeOrDed = 1, + is_causal: Optional[CausalArg3DTypeOrDed] = False, + scale: Optional[float] = None, + q_tile_shape: Optional[Dimension3DType] = None, + kv_tile_shape: Optional[Dimension3DType] = None, + torch_compile: bool = False, + return_lse: bool = False, +) -> Union[Tensor, Tuple[Tensor, Tensor]]: + return flex_fna_generic( + query=query, + key=key, + value=value, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + scale=scale, + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + torch_compile=torch_compile, + return_lse=return_lse, + ) diff --git a/natten/torch-ext/natten/backends/fmha.py b/natten/torch-ext/natten/backends/fmha.py new file mode 100644 index 00000000..139f0985 --- /dev/null +++ b/natten/torch-ext/natten/backends/fmha.py @@ -0,0 +1,283 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# +import functools +from typing import Optional, Tuple, Union + +import torch +from torch import Tensor +from torch.amp import custom_bwd, custom_fwd +from torch.autograd import Function + +amp_fwd = functools.partial(custom_fwd, device_type="cuda") +amp_bwd = functools.partial(custom_bwd, device_type="cuda") + +from .._libnatten import fmha_backward, fmha_forward +from ..backends.configs.checks import can_run_cutlass_fmha +from ..backends.configs.cutlass import ( + check_cutlass_fmha_backward_config, + check_cutlass_fmha_forward_config, +) +from ..types import ( + CutlassFmhaBackwardConfigType, + CutlassFmhaForwardConfigType, + NoneType, +) +from ..utils import log +from ..utils.checks import fmha_tensor_checks, varlen_tensor_checks + +logger = log.get_logger(__name__) + + +class CutlassFmhaAutogradFn(Function): + @staticmethod + @amp_fwd + def forward( + ctx, + query: Tensor, + key: Tensor, + value: Tensor, + is_causal: bool, + scale: float, + forward_config: CutlassFmhaForwardConfigType, + backward_config: CutlassFmhaBackwardConfigType, + backward_kv_splits: Optional[int], + backward_use_pt_reduction: bool, + cumulative_seqlen_Q: Optional[Tensor], + cumulative_seqlen_KV: Optional[Tensor], + max_seqlen_Q: int, + max_seqlen_KV: int, + ) -> Tuple[Tensor, Tensor]: + + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + + assert query.dim() == value.dim() == 4 + assert query.shape[0] == value.shape[0] + assert query.shape[2] == value.shape[2] + + q_tile_size, kv_tile_size = forward_config + output, logsumexp = fmha_forward( + query, + key, + value, + is_causal, + scale, + q_tile_size, + kv_tile_size, + cumulative_seqlen_Q, + cumulative_seqlen_KV, + max_seqlen_Q, + max_seqlen_KV, + ) + + ctx.save_for_backward( + query, + key, + value, + logsumexp, + output, + cumulative_seqlen_Q, + cumulative_seqlen_KV, + ) + ctx.scale = scale + ctx.is_causal = is_causal + ctx.max_seqlen_Q = max_seqlen_Q + ctx.max_seqlen_KV = max_seqlen_KV + ctx.backward_config = backward_config + ctx.backward_kv_splits = backward_kv_splits + ctx.backward_use_pt_reduction = backward_use_pt_reduction + # Always record determinism behavior during forward pass (forward pass itself is + # deterministic anyway). + # Determinism could be limited to part of the program, which means during forward pass + # it'll be true, but on .backward() call, if it's been turned off, it will stay off when we + # get to this operation's backward call. + ctx.deterministic = torch.are_deterministic_algorithms_enabled() + + return output, logsumexp + + @staticmethod + @amp_bwd + def backward(ctx, grad_out: Tensor, grad_lse: Tensor) -> Tuple[ + Tensor, + Tensor, + Tensor, + NoneType, + NoneType, + NoneType, + NoneType, + # kv_splits, use_pt_reduction + NoneType, + NoneType, + # varlen + NoneType, + NoneType, + NoneType, + NoneType, + ]: + ( + query, + key, + value, + logsumexp, + output, + cumulative_seqlen_Q, + cumulative_seqlen_KV, + ) = ctx.saved_tensors + d_output = grad_out.contiguous() + + q_tile_size, k_tile_size = ctx.backward_config + + d_query, d_key, d_value = fmha_backward( + query, + key, + value, + output, + d_output, + logsumexp, + ctx.is_causal, + ctx.scale, + q_tile_size, + k_tile_size, + ctx.backward_kv_splits, + ctx.backward_use_pt_reduction, + cumulative_seqlen_Q, + cumulative_seqlen_KV, + ctx.max_seqlen_Q, + ctx.max_seqlen_KV, + ctx.deterministic, + ) + + return ( + d_query, + d_key, + d_value, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + ) + + +def cutlass_fmha( + query: Tensor, + key: Tensor, + value: Tensor, + is_causal: bool = False, + scale: Optional[float] = None, + q_tile_size: Optional[int] = None, + kv_tile_size: Optional[int] = None, + backward_q_tile_size: Optional[int] = None, + backward_kv_tile_size: Optional[int] = None, + backward_kv_splits: Optional[int] = None, + backward_use_pt_reduction: bool = False, + return_lse: bool = False, + # varlen parameters + cumulative_seqlen_Q: Optional[Tensor] = None, + cumulative_seqlen_KV: Optional[Tensor] = None, + max_seqlen_Q: int = 0, + max_seqlen_KV: int = 0, +) -> Union[Tensor, Tuple[Tensor, Tensor]]: + + fmha_tensor_checks( + query, + key, + value, + must_match_head_dims=False, + supports_gqa_mqa=True, + backend_name="CUTLASS FMHA", + ) + + ( + cumulative_seqlen_Q, + cumulative_seqlen_KV, + max_seqlen_Q, + max_seqlen_KV, + ) = varlen_tensor_checks( + query=query, + key=key, + value=value, + cumulative_seqlen_Q=cumulative_seqlen_Q, + cumulative_seqlen_KV=cumulative_seqlen_KV, + max_seqlen_Q=max_seqlen_Q, + max_seqlen_KV=max_seqlen_KV, + ) + is_varlen = cumulative_seqlen_Q is not None + + assert can_run_cutlass_fmha( + query, key, value, is_causal=is_causal, is_varlen=is_varlen, raise_error=True + ) + + forward_config = check_cutlass_fmha_forward_config( + input_tensor=query if value.shape[-1] <= query.shape[-1] else value, + q_tile_size=q_tile_size, + kv_tile_size=kv_tile_size, + ) + + backward_config = check_cutlass_fmha_backward_config( + input_tensor=key if value.shape[-1] <= key.shape[-1] else value, + q_tile_size=backward_q_tile_size, + kv_tile_size=backward_kv_tile_size, + ) + + scale = scale or query.shape[-1] ** -0.5 + + # GQA/MQA is not supported by the kernel; only allowed via graph transform + is_gqa = query.shape[-2] != key.shape[-2] + if is_gqa: + heads = query.shape[-2] + heads_kv = key.shape[-2] + assert key.shape[-2] == value.shape[-2] + assert heads >= heads_kv + assert heads % heads_kv == 0 + h_k = heads // heads_kv + + key = torch.repeat_interleave(key, repeats=h_k, dim=-2, output_size=heads) + value = torch.repeat_interleave(value, repeats=h_k, dim=-2, output_size=heads) + + output, lse = CutlassFmhaAutogradFn.apply( + query, + key, + value, + is_causal, + scale, + forward_config, + backward_config, + backward_kv_splits, + backward_use_pt_reduction, + cumulative_seqlen_Q, + cumulative_seqlen_KV, + max_seqlen_Q, + max_seqlen_KV, + ) + + if return_lse: + return output, lse + + return output diff --git a/natten/torch-ext/natten/backends/fna.py b/natten/torch-ext/natten/backends/fna.py new file mode 100644 index 00000000..15be3930 --- /dev/null +++ b/natten/torch-ext/natten/backends/fna.py @@ -0,0 +1,417 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# +import functools +from typing import Optional, Tuple, Union + +import torch +from torch import Tensor +from torch.amp import custom_bwd, custom_fwd +from torch.autograd import Function + +amp_fwd = functools.partial(custom_fwd, device_type="cuda") +amp_bwd = functools.partial(custom_bwd, device_type="cuda") + +from .._libnatten import ( + na1d_backward, + na1d_forward, + na2d_backward, + na2d_forward, + na3d_backward, + na3d_forward, +) +from ..backends.configs.checks import can_run_cutlass_fna +from ..backends.configs.cutlass import ( + check_cutlass_fna_backward_config, + check_cutlass_fna_forward_config, +) +from ..types import ( + CausalArg1DTypeOrDed, + CausalArg2DTypeOrDed, + CausalArg3DTypeOrDed, + CausalArgType, + CausalArgTypeOrDed, + CutlassFnaBackwardConfigType, + CutlassFnaForwardConfigType, + Dimension1DType, + Dimension1DTypeOrDed, + Dimension2DType, + Dimension2DTypeOrDed, + Dimension3DType, + Dimension3DTypeOrDed, + DimensionType, + DimensionTypeOrDed, + NoneType, +) +from ..utils import log +from ..utils.checks import ( + check_all_args, + check_args_against_input, + na_tensor_checks, +) + +logger = log.get_logger(__name__) + + +def make_cutlass_fna_autograd_fn(na_dim): + assert na_dim in [1, 2, 3] + + FORWARD_OPS = { + 1: na1d_forward, + 2: na2d_forward, + 3: na3d_forward, + } + + BACKWARD_OPS = { + 1: na1d_backward, + 2: na2d_backward, + 3: na3d_backward, + } + + class CutlassFnaGenericAutogradFn(Function): + @staticmethod + @amp_fwd + def forward( + ctx, + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size: DimensionType, + stride: DimensionType, + dilation: DimensionType, + is_causal: CausalArgType, + scale: float, + forward_config: CutlassFnaForwardConfigType, + backward_config: CutlassFnaBackwardConfigType, + backward_kv_splits: Optional[DimensionType], + backward_use_pt_reduction: bool, + ) -> Tuple[Tensor, Tensor]: + kernel_size, stride, dilation, is_causal = check_all_args( + na_dim, kernel_size, stride, dilation, is_causal + ) + + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + + assert query.dim() == value.dim() == 3 + na_dim + assert query.shape[0] == value.shape[0] + assert query.shape[-2] == value.shape[-2] + + q_tile_shape, kv_tile_shape = forward_config + + output, logsumexp = FORWARD_OPS[na_dim]( + query, + key, + value, + kernel_size, + stride, + dilation, + is_causal, + scale, + q_tile_shape, + kv_tile_shape, + ) + + ctx.save_for_backward(query, key, value, logsumexp, output) + ctx.kernel_size = kernel_size + ctx.stride = stride + ctx.dilation = dilation + ctx.is_causal = is_causal + ctx.scale = scale + ctx.backward_config = backward_config + ctx.backward_kv_splits = backward_kv_splits + ctx.backward_use_pt_reduction = backward_use_pt_reduction + # Always record determinism behavior during forward pass (forward pass itself is + # deterministic anyway). + # Determinism could be limited to part of the program, which means during forward pass + # it'll be true, but on .backward() call, if it's been turned off, it will stay off when we + # get to this operation's backward call. + ctx.deterministic = torch.are_deterministic_algorithms_enabled() + + return output, logsumexp + + @staticmethod + @amp_bwd + def backward(ctx, grad_out: Tensor, grad_lse: Tensor) -> Tuple[ + Tensor, + Tensor, + Tensor, + NoneType, + NoneType, + NoneType, + NoneType, + NoneType, + NoneType, + NoneType, + NoneType, + NoneType, + ]: + query, key, value, logsumexp, output = ctx.saved_tensors + d_output = grad_out.contiguous() + + q_tile_shape, k_tile_shape = ctx.backward_config + + d_query, d_key, d_value = BACKWARD_OPS[na_dim]( + query, + key, + value, + output, + d_output, + logsumexp, + ctx.kernel_size, + ctx.stride, + ctx.dilation, + ctx.is_causal, + ctx.scale, + q_tile_shape, + k_tile_shape, + ctx.backward_kv_splits, + ctx.backward_use_pt_reduction, + ctx.deterministic, + ) + + return ( + d_query, + d_key, + d_value, + None, + None, + None, + None, + None, + None, + None, + None, + None, + ) + + return CutlassFnaGenericAutogradFn + + +CutlassFna1DAutogradFn = make_cutlass_fna_autograd_fn(1) +CutlassFna2DAutogradFn = make_cutlass_fna_autograd_fn(2) +CutlassFna3DAutogradFn = make_cutlass_fna_autograd_fn(3) + + +CutlassFNAAutogradFns = { + 1: CutlassFna1DAutogradFn, + 2: CutlassFna2DAutogradFn, + 3: CutlassFna3DAutogradFn, +} + + +def cutlass_fna_generic( + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size: DimensionTypeOrDed, + stride: DimensionTypeOrDed = 1, + dilation: DimensionTypeOrDed = 1, + is_causal: Optional[CausalArgTypeOrDed] = False, + scale: Optional[float] = None, + q_tile_shape: Optional[DimensionType] = None, + kv_tile_shape: Optional[DimensionType] = None, + backward_q_tile_shape: Optional[DimensionType] = None, + backward_kv_tile_shape: Optional[DimensionType] = None, + backward_kv_splits: Optional[DimensionType] = None, + backward_use_pt_reduction: bool = False, + return_lse: bool = False, +) -> Union[Tensor, Tuple[Tensor, Tensor]]: + + na_tensor_checks( + query, key, value, must_match_head_dims=False, supports_gqa_mqa=True + ) + + assert can_run_cutlass_fna(query, key, value, raise_error=True) + + na_dim = query.dim() - 3 # batch, heads, head_dim + + assert na_dim in [1, 2, 3] + + kernel_size, stride, dilation, is_causal = check_all_args( + na_dim, kernel_size, stride, dilation, is_causal + ) + + check_args_against_input( + query, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + ) + + forward_config = check_cutlass_fna_forward_config( + input_tensor=query if value.shape[-1] <= query.shape[-1] else value, + dilation=dilation, + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + ) + + backward_config = check_cutlass_fna_backward_config( + input_tensor=key if value.shape[-1] <= key.shape[-1] else value, + q_tile_shape=backward_q_tile_shape, + kv_tile_shape=backward_kv_tile_shape, + ) + + scale = scale or query.shape[-1] ** -0.5 + + # GQA/MQA is not supported by the kernel; only allowed via graph transform + is_gqa = query.shape[-2] != key.shape[-2] + if is_gqa: + heads = query.shape[-2] + heads_kv = key.shape[-2] + assert key.shape[-2] == value.shape[-2] + assert heads >= heads_kv + assert heads % heads_kv == 0 + h_k = heads // heads_kv + + key = torch.repeat_interleave(key, repeats=h_k, dim=-2, output_size=heads) + value = torch.repeat_interleave(value, repeats=h_k, dim=-2, output_size=heads) + + output, lse = CutlassFNAAutogradFns[na_dim].apply( + query, + key, + value, + kernel_size, + stride, + dilation, + is_causal, + scale, + forward_config, + backward_config, + backward_kv_splits, + backward_use_pt_reduction, + ) + + if return_lse: + return output, lse + + return output + + +def na1d_cutlass_fna( + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size: Dimension1DTypeOrDed, + stride: Dimension1DTypeOrDed = 1, + dilation: Dimension1DTypeOrDed = 1, + is_causal: Optional[CausalArg1DTypeOrDed] = False, + scale: Optional[float] = None, + q_tile_shape: Optional[Dimension1DType] = None, + kv_tile_shape: Optional[Dimension1DType] = None, + backward_q_tile_shape: Optional[Dimension1DType] = None, + backward_kv_tile_shape: Optional[Dimension1DType] = None, + backward_kv_splits: Optional[Dimension1DType] = None, + backward_use_pt_reduction: bool = False, + return_lse: bool = False, +) -> Union[Tensor, Tuple[Tensor, Tensor]]: + return cutlass_fna_generic( + query=query, + key=key, + value=value, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + scale=scale, + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + backward_q_tile_shape=backward_q_tile_shape, + backward_kv_tile_shape=backward_kv_tile_shape, + backward_kv_splits=backward_kv_splits, + backward_use_pt_reduction=backward_use_pt_reduction, + return_lse=return_lse, + ) + + +def na2d_cutlass_fna( + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size: Dimension2DTypeOrDed, + stride: Dimension2DTypeOrDed = 1, + dilation: Dimension2DTypeOrDed = 1, + is_causal: Optional[CausalArg2DTypeOrDed] = False, + scale: Optional[float] = None, + q_tile_shape: Optional[Dimension2DType] = None, + kv_tile_shape: Optional[Dimension2DType] = None, + backward_q_tile_shape: Optional[Dimension2DType] = None, + backward_kv_tile_shape: Optional[Dimension2DType] = None, + backward_kv_splits: Optional[Dimension2DType] = None, + backward_use_pt_reduction: bool = False, + return_lse: bool = False, +) -> Union[Tensor, Tuple[Tensor, Tensor]]: + return cutlass_fna_generic( + query=query, + key=key, + value=value, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + scale=scale, + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + backward_q_tile_shape=backward_q_tile_shape, + backward_kv_tile_shape=backward_kv_tile_shape, + backward_kv_splits=backward_kv_splits, + backward_use_pt_reduction=backward_use_pt_reduction, + return_lse=return_lse, + ) + + +def na3d_cutlass_fna( + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size: Dimension3DTypeOrDed, + stride: Dimension3DTypeOrDed = 1, + dilation: Dimension3DTypeOrDed = 1, + is_causal: Optional[CausalArg3DTypeOrDed] = False, + scale: Optional[float] = None, + q_tile_shape: Optional[Dimension3DType] = None, + kv_tile_shape: Optional[Dimension3DType] = None, + backward_q_tile_shape: Optional[Dimension3DType] = None, + backward_kv_tile_shape: Optional[Dimension3DType] = None, + backward_kv_splits: Optional[Dimension3DType] = None, + backward_use_pt_reduction: bool = False, + return_lse: bool = False, +) -> Union[Tensor, Tuple[Tensor, Tensor]]: + return cutlass_fna_generic( + query=query, + key=key, + value=value, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + scale=scale, + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + backward_q_tile_shape=backward_q_tile_shape, + backward_kv_tile_shape=backward_kv_tile_shape, + backward_kv_splits=backward_kv_splits, + backward_use_pt_reduction=backward_use_pt_reduction, + return_lse=return_lse, + ) diff --git a/natten/torch-ext/natten/backends/hopper_fmha.py b/natten/torch-ext/natten/backends/hopper_fmha.py new file mode 100644 index 00000000..4f98a51e --- /dev/null +++ b/natten/torch-ext/natten/backends/hopper_fmha.py @@ -0,0 +1,261 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# +import functools +from typing import Optional, Tuple, Union + +import torch +from torch import Tensor +from torch.amp import custom_bwd, custom_fwd +from torch.autograd import Function + +amp_fwd = functools.partial(custom_fwd, device_type="cuda") +amp_bwd = functools.partial(custom_bwd, device_type="cuda") + +from .._libnatten import hopper_fmha_backward, hopper_fmha_forward +from ..backends.configs.checks import can_run_cutlass_hopper_fmha +from ..backends.configs.cutlass_hopper import ( + check_cutlass_hopper_fmha_backward_config, + check_cutlass_hopper_fmha_forward_config, +) +from ..types import ( + CutlassHopperFmhaBackwardConfigType, + CutlassHopperFmhaForwardConfigType, + KernelSchedule, + NoneType, +) +from ..utils import log +from ..utils.checks import fmha_tensor_checks, varlen_tensor_checks + +logger = log.get_logger(__name__) + + +class CutlassHopperFmhaAutogradFn(Function): + @staticmethod + @amp_fwd + def forward( + ctx, + query: Tensor, + key: Tensor, + value: Tensor, + is_causal: bool, + scale: float, + forward_config: CutlassHopperFmhaForwardConfigType, + backward_config: CutlassHopperFmhaBackwardConfigType, + cumulative_seqlen_Q: Optional[Tensor], + cumulative_seqlen_KV: Optional[Tensor], + max_seqlen_Q: int, + max_seqlen_KV: int, + ) -> Tuple[Tensor, Tensor]: + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + + (q_tile_size, kv_tile_size), kernel_schedule = forward_config + + output, logsumexp = hopper_fmha_forward( + query, + key, + value, + is_causal, + scale, + q_tile_size, + kv_tile_size, + kernel_schedule.value, # TODO: I don't like this -- write a map with checks? + cumulative_seqlen_Q, + cumulative_seqlen_KV, + max_seqlen_Q, + max_seqlen_KV, + ) + + ctx.save_for_backward( + query, + key, + value, + logsumexp, + output, + cumulative_seqlen_Q, + cumulative_seqlen_KV, + ) + ctx.scale = scale + ctx.is_causal = is_causal + ctx.max_seqlen_Q = max_seqlen_Q + ctx.max_seqlen_KV = max_seqlen_KV + ctx.backward_config = backward_config + # Always record determinism behavior during forward pass (forward pass itself is + # deterministic anyway). + # Determinism could be limited to part of the program, which means during forward pass + # it'll be true, but on .backward() call, if it's been turned off, it will stay off when we + # get to this operation's backward call. + ctx.deterministic = torch.are_deterministic_algorithms_enabled() + + return output, logsumexp + + @staticmethod + @amp_bwd + def backward(ctx, grad_out: Tensor, grad_lse: Tensor) -> Tuple[ + Tensor, + Tensor, + Tensor, + NoneType, + NoneType, + NoneType, + NoneType, + NoneType, + NoneType, + NoneType, + NoneType, + ]: + ( + query, + key, + value, + logsumexp, + output, + cumulative_seqlen_Q, + cumulative_seqlen_KV, + ) = ctx.saved_tensors + d_output = grad_out.contiguous() # noqa: F841 + + q_tile_size, k_tile_size = ctx.backward_config + + if ctx.deterministic: + raise RuntimeError( + "Hopper FMHA backward pass does not have a deterministic mode, " + "but PyTorch's deterministic algorithms were enabled. To proceed, " + "you must either disable torch's deterministic mode, or choose a " + "different backend." + ) + + d_query, d_key, d_value = hopper_fmha_backward( + query, + key, + value, + output, + d_output, + logsumexp, + ctx.is_causal, + ctx.scale, + q_tile_size, + k_tile_size, + cumulative_seqlen_Q, + cumulative_seqlen_KV, + ctx.max_seqlen_Q, + ctx.max_seqlen_KV, + ) + + return d_query, d_key, d_value, None, None, None, None, None, None, None, None + + +def cutlass_hopper_fmha( + query: Tensor, + key: Tensor, + value: Tensor, + is_causal: bool = False, + scale: Optional[float] = None, + q_tile_size: Optional[int] = None, + kv_tile_size: Optional[int] = None, + kernel_schedule: Optional[KernelSchedule] = None, + backward_q_tile_size: Optional[int] = None, + backward_kv_tile_size: Optional[int] = None, + return_lse: bool = False, + # varlen parameters + cumulative_seqlen_Q: Optional[Tensor] = None, + cumulative_seqlen_KV: Optional[Tensor] = None, + max_seqlen_Q: int = 0, + max_seqlen_KV: int = 0, +) -> Union[Tensor, Tuple[Tensor, Tensor]]: + + fmha_tensor_checks( + query, + key, + value, + must_match_head_dims=True, + supports_gqa_mqa=True, + backend_name="Hopper FMHA", + ) + + ( + cumulative_seqlen_Q, + cumulative_seqlen_KV, + max_seqlen_Q, + max_seqlen_KV, + ) = varlen_tensor_checks( + query=query, + key=key, + value=value, + cumulative_seqlen_Q=cumulative_seqlen_Q, + cumulative_seqlen_KV=cumulative_seqlen_KV, + max_seqlen_Q=max_seqlen_Q, + max_seqlen_KV=max_seqlen_KV, + ) + is_varlen = cumulative_seqlen_Q is not None + + assert can_run_cutlass_hopper_fmha( + query, key, value, is_causal=is_causal, is_varlen=is_varlen, raise_error=True + ) + + forward_config = check_cutlass_hopper_fmha_forward_config( + input_tensor=query, + q_tile_size=q_tile_size, + kv_tile_size=kv_tile_size, + kernel_schedule=kernel_schedule, + ) + backward_config = check_cutlass_hopper_fmha_backward_config( + input_tensor=query, + q_tile_size=backward_q_tile_size, + kv_tile_size=backward_kv_tile_size, + ) + + scale = scale or query.shape[-1] ** -0.5 + + # GQA/MQA is not supported by the kernel; only allowed via graph transform + is_gqa = query.shape[-2] != key.shape[-2] + if is_gqa: + heads = query.shape[-2] + heads_kv = key.shape[-2] + assert key.shape[-2] == value.shape[-2] + assert heads >= heads_kv + assert heads % heads_kv == 0 + h_k = heads // heads_kv + + key = torch.repeat_interleave(key, repeats=h_k, dim=-2, output_size=heads) + value = torch.repeat_interleave(value, repeats=h_k, dim=-2, output_size=heads) + + output, lse = CutlassHopperFmhaAutogradFn.apply( + query, + key, + value, + is_causal, + scale, + forward_config, + backward_config, + cumulative_seqlen_Q, + cumulative_seqlen_KV, + max_seqlen_Q, + max_seqlen_KV, + ) + + if return_lse: + return output, lse + + return output diff --git a/natten/torch-ext/natten/backends/hopper_fna.py b/natten/torch-ext/natten/backends/hopper_fna.py new file mode 100644 index 00000000..dff38175 --- /dev/null +++ b/natten/torch-ext/natten/backends/hopper_fna.py @@ -0,0 +1,512 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# +import functools +from typing import Optional, Tuple, Union + +import torch +from torch import Tensor +from torch.amp import custom_bwd, custom_fwd +from torch.autograd import Function + +amp_fwd = functools.partial(custom_fwd, device_type="cuda") +amp_bwd = functools.partial(custom_bwd, device_type="cuda") + +from .._libnatten import ( + hopper_na1d_backward, + hopper_na1d_forward, + hopper_na2d_backward, + hopper_na2d_forward, + hopper_na3d_backward, + hopper_na3d_forward, +) +from ..backends.configs.checks import can_run_cutlass_hopper_fna +from ..backends.configs.cutlass_hopper import ( + check_cutlass_hopper_fna_backward_config, + check_cutlass_hopper_fna_forward_config, +) +from ..token_permute import token_permute_operation, token_unpermute_operation +from ..types import ( + CausalArg1DTypeOrDed, + CausalArg2DTypeOrDed, + CausalArg3DTypeOrDed, + CausalArgType, + CausalArgTypeOrDed, + CutlassHopperFnaBackwardConfigType, + CutlassHopperFnaForwardConfigType, + Dimension1DType, + Dimension1DTypeOrDed, + Dimension2DType, + Dimension2DTypeOrDed, + Dimension3DType, + Dimension3DTypeOrDed, + DimensionType, + DimensionTypeOrDed, + KernelSchedule, + NoneType, +) +from ..utils.checks import ( + check_all_args, + check_args_against_input, + na_tensor_checks, +) + + +def make_cutlass_hopper_fna_autograd_fn(na_dim): + assert na_dim in [1, 2, 3] + + FORWARD_OPS = { + 1: hopper_na1d_forward, + 2: hopper_na2d_forward, + 3: hopper_na3d_forward, + } + + BACKWARD_OPS = { + 1: hopper_na1d_backward, + 2: hopper_na2d_backward, + 3: hopper_na3d_backward, + } + + class CutlassHopperFnaGenericAutogradFn(Function): + @staticmethod + @amp_fwd + def forward( + ctx, + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size: DimensionType, + stride: DimensionType, + dilation: DimensionType, + is_causal: CausalArgType, + scale: float, + forward_config: CutlassHopperFnaForwardConfigType, + backward_config: CutlassHopperFnaBackwardConfigType, + ) -> Tuple[Tensor, Tensor]: + kernel_size, stride, dilation, is_causal = check_all_args( + na_dim, kernel_size, stride, dilation, is_causal + ) + + (q_tile_shape, kv_tile_shape), kernel_schedule = forward_config + + # Token permute begin + query_perm, qkv_shape, q_shape = token_permute_operation( + query, + q_tile_shape, + dilation=dilation, + flip_tiled_dims=True, + ) + key_perm, _, k_shape = token_permute_operation( + key, kv_tile_shape, dilation=dilation, flip_tiled_dims=True + ) + value_perm, _, v_shape = token_permute_operation( + value, kv_tile_shape, dilation=dilation, flip_tiled_dims=True + ) + + assert k_shape == v_shape + kv_shape = k_shape + # Token permute end + + query_perm = query_perm.contiguous() + key_perm = key_perm.contiguous() + value_perm = value_perm.contiguous() + + output_perm, logsumexp_perm = FORWARD_OPS[na_dim]( + query_perm, + key_perm, + value_perm, + kernel_size, + stride, + dilation, + is_causal, + scale, + q_shape, + kv_shape, + qkv_shape, + q_tile_shape, + kv_tile_shape, + kernel_schedule.value, # TODO: I don't like this -- write a map with checks? + ) + + # Token un-permute begin + output = token_unpermute_operation( + output_perm, + token_layout_shape=qkv_shape, + tile_shape=q_tile_shape, + dilation=dilation, + flip_tiled_dims=True, + ) + logsumexp = token_unpermute_operation( + logsumexp_perm.unsqueeze(-1), + token_layout_shape=qkv_shape, + tile_shape=q_tile_shape, + dilation=dilation, + flip_tiled_dims=True, + ).squeeze(-1) + # Token un-permute end + + ctx.save_for_backward(query, key, value, logsumexp, output) + ctx.kernel_size = kernel_size + ctx.stride = stride + ctx.dilation = dilation + ctx.is_causal = is_causal + ctx.scale = scale + ctx.backward_config = backward_config + # Always record determinism behavior during forward pass (forward pass itself is + # deterministic anyway). + # Determinism could be limited to part of the program, which means during forward pass + # it'll be true, but on .backward() call, if it's been turned off, it will stay off when we + # get to this operation's backward call. + ctx.deterministic = torch.are_deterministic_algorithms_enabled() + + return output, logsumexp + + @staticmethod + @amp_bwd + def backward(ctx, d_output: Tensor, d_lse: Tensor) -> Tuple[ + Tensor, + Tensor, + Tensor, + NoneType, + NoneType, + NoneType, + NoneType, + NoneType, + NoneType, + NoneType, + ]: + query, key, value, logsumexp, output = ctx.saved_tensors + kernel_size, stride, dilation, is_causal, scale = ( + ctx.kernel_size, + ctx.stride, + ctx.dilation, + ctx.is_causal, + ctx.scale, + ) + + q_tile_shape, kv_tile_shape = ctx.backward_config + + if ctx.deterministic: + raise RuntimeError( + "Hopper FNA backward pass does not have a deterministic mode, " + "but PyTorch's deterministic algorithms were enabled. To proceed, " + "you must either disable torch's deterministic mode, or choose a " + "different backend." + ) + + # Token permute begin + + query_perm, qkv_shape, q_shape = token_permute_operation( + query, tile_shape=q_tile_shape, dilation=dilation, flip_tiled_dims=True + ) + output_perm, _, o_shape = token_permute_operation( + output, tile_shape=q_tile_shape, dilation=dilation, flip_tiled_dims=True + ) + d_output_perm, _, d_o_shape = token_permute_operation( + d_output, + tile_shape=q_tile_shape, + dilation=dilation, + flip_tiled_dims=True, + ) + logsumexp_perm, _, _ = token_permute_operation( + logsumexp.unsqueeze(-1), + tile_shape=q_tile_shape, + dilation=dilation, + flip_tiled_dims=True, + ) + key_perm, _, k_shape = token_permute_operation( + key, tile_shape=kv_tile_shape, dilation=dilation, flip_tiled_dims=True + ) + value_perm, _, v_shape = token_permute_operation( + value, tile_shape=kv_tile_shape, dilation=dilation, flip_tiled_dims=True + ) + + assert q_shape == o_shape == d_o_shape + assert k_shape == v_shape + kv_shape = k_shape + # Token permute end + + query_perm = query_perm.contiguous() + key_perm = key_perm.contiguous() + value_perm = value_perm.contiguous() + output_perm = output_perm.contiguous() + d_output_perm = d_output_perm.contiguous() + logsumexp_perm = logsumexp_perm.squeeze(-1) + + d_query_perm, d_key_perm, d_value_perm = BACKWARD_OPS[na_dim]( + query_perm, + key_perm, + value_perm, + output_perm, + d_output_perm, + logsumexp_perm, + kernel_size, + stride, + dilation, + is_causal, + scale, + q_shape, + kv_shape, + qkv_shape, + q_tile_shape, + kv_tile_shape, + ) + + # Token un-permute begin + d_query = token_unpermute_operation( + d_query_perm, + token_layout_shape=qkv_shape, + tile_shape=q_tile_shape, + dilation=dilation, + flip_tiled_dims=True, + ) + d_key = token_unpermute_operation( + d_key_perm, + token_layout_shape=qkv_shape, + tile_shape=kv_tile_shape, + dilation=dilation, + flip_tiled_dims=True, + ) + d_value = token_unpermute_operation( + d_value_perm, + token_layout_shape=qkv_shape, + tile_shape=kv_tile_shape, + dilation=dilation, + flip_tiled_dims=True, + ) + # Token un-permute end + + assert d_query.shape == query.shape + assert d_key.shape == key.shape + assert d_value.shape == value.shape + + return ( + d_query, + d_key, + d_value, + None, + None, + None, + None, + None, + None, + None, + ) + + return CutlassHopperFnaGenericAutogradFn + + +CutlassHopperFna1DAutogradFn = make_cutlass_hopper_fna_autograd_fn(1) +CutlassHopperFna2DAutogradFn = make_cutlass_hopper_fna_autograd_fn(2) +CutlassHopperFna3DAutogradFn = make_cutlass_hopper_fna_autograd_fn(3) + + +CutlassHopperFNAAutogradFns = { + 1: CutlassHopperFna1DAutogradFn, + 2: CutlassHopperFna2DAutogradFn, + 3: CutlassHopperFna3DAutogradFn, +} + + +def cutlass_hopper_fna_generic( + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size: DimensionTypeOrDed, + stride: DimensionTypeOrDed = 1, + dilation: DimensionTypeOrDed = 1, + is_causal: Optional[CausalArgTypeOrDed] = False, + scale: Optional[float] = None, + q_tile_shape: Optional[DimensionType] = None, + kv_tile_shape: Optional[DimensionType] = None, + backward_q_tile_shape: Optional[DimensionType] = None, + backward_kv_tile_shape: Optional[DimensionType] = None, + kernel_schedule: Optional[KernelSchedule] = None, + return_lse: bool = False, +) -> Union[Tensor, Tuple[Tensor, Tensor]]: + + na_tensor_checks( + query, key, value, must_match_head_dims=True, supports_gqa_mqa=True + ) + + assert can_run_cutlass_hopper_fna(query, key, value, raise_error=True) + + na_dim = query.dim() - 3 # batch, heads, head_dim + + kernel_size, stride, dilation, is_causal = check_all_args( + na_dim, kernel_size, stride, dilation, is_causal + ) + + check_args_against_input( + query, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + ) + + forward_config = check_cutlass_hopper_fna_forward_config( + input_tensor=query, + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + kernel_schedule=kernel_schedule, + ) + backward_config = check_cutlass_hopper_fna_backward_config( + input_tensor=query, + q_tile_shape=backward_q_tile_shape, + kv_tile_shape=backward_kv_tile_shape, + ) + + scale = scale or query.shape[-1] ** -0.5 + + # GQA/MQA is not supported by the kernel; only allowed via graph transform + is_gqa = query.shape[-2] != key.shape[-2] + if is_gqa: + heads = query.shape[-2] + heads_kv = key.shape[-2] + assert key.shape[-2] == value.shape[-2] + assert heads >= heads_kv + assert heads % heads_kv == 0 + h_k = heads // heads_kv + + key = torch.repeat_interleave(key, repeats=h_k, dim=-2, output_size=heads) + value = torch.repeat_interleave(value, repeats=h_k, dim=-2, output_size=heads) + + output, lse = CutlassHopperFNAAutogradFns[na_dim].apply( + query, + key, + value, + kernel_size, + stride, + dilation, + is_causal, + scale, + forward_config, + backward_config, + ) + + if return_lse: + return output, lse + + return output + + +def na1d_cutlass_hopper_fna( + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size: Dimension1DTypeOrDed, + stride: Dimension1DTypeOrDed = 1, + dilation: Dimension1DTypeOrDed = 1, + is_causal: Optional[CausalArg1DTypeOrDed] = False, + scale: Optional[float] = None, + q_tile_shape: Optional[Dimension1DType] = None, + kv_tile_shape: Optional[Dimension1DType] = None, + backward_q_tile_shape: Optional[Dimension1DType] = None, + backward_kv_tile_shape: Optional[Dimension1DType] = None, + kernel_schedule: Optional[KernelSchedule] = None, + return_lse: bool = False, +) -> Union[Tensor, Tuple[Tensor, Tensor]]: + return cutlass_hopper_fna_generic( + query=query, + key=key, + value=value, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + scale=scale, + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + backward_q_tile_shape=backward_q_tile_shape, + backward_kv_tile_shape=backward_kv_tile_shape, + kernel_schedule=kernel_schedule, + return_lse=return_lse, + ) + + +def na2d_cutlass_hopper_fna( + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size: Dimension2DTypeOrDed, + stride: Dimension2DTypeOrDed = 1, + dilation: Dimension2DTypeOrDed = 1, + is_causal: Optional[CausalArg2DTypeOrDed] = False, + scale: Optional[float] = None, + q_tile_shape: Optional[Dimension2DType] = None, + kv_tile_shape: Optional[Dimension2DType] = None, + backward_q_tile_shape: Optional[Dimension2DType] = None, + backward_kv_tile_shape: Optional[Dimension2DType] = None, + kernel_schedule: Optional[KernelSchedule] = None, + return_lse: bool = False, +) -> Union[Tensor, Tuple[Tensor, Tensor]]: + return cutlass_hopper_fna_generic( + query=query, + key=key, + value=value, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + scale=scale, + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + backward_q_tile_shape=backward_q_tile_shape, + backward_kv_tile_shape=backward_kv_tile_shape, + kernel_schedule=kernel_schedule, + return_lse=return_lse, + ) + + +def na3d_cutlass_hopper_fna( + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size: Dimension3DTypeOrDed, + stride: Dimension3DTypeOrDed = 1, + dilation: Dimension3DTypeOrDed = 1, + is_causal: Optional[CausalArg3DTypeOrDed] = False, + scale: Optional[float] = None, + q_tile_shape: Optional[Dimension3DType] = None, + kv_tile_shape: Optional[Dimension3DType] = None, + backward_q_tile_shape: Optional[Dimension3DType] = None, + backward_kv_tile_shape: Optional[Dimension3DType] = None, + kernel_schedule: Optional[KernelSchedule] = None, + return_lse: bool = False, +) -> Union[Tensor, Tuple[Tensor, Tensor]]: + return cutlass_hopper_fna_generic( + query=query, + key=key, + value=value, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + scale=scale, + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + backward_q_tile_shape=backward_q_tile_shape, + backward_kv_tile_shape=backward_kv_tile_shape, + kernel_schedule=kernel_schedule, + return_lse=return_lse, + ) diff --git a/natten/torch-ext/natten/backends/reference.py b/natten/torch-ext/natten/backends/reference.py new file mode 100644 index 00000000..c6f72d74 --- /dev/null +++ b/natten/torch-ext/natten/backends/reference.py @@ -0,0 +1,343 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# +import functools +from typing import Optional, Tuple, Union + +import torch +from torch import Tensor +from torch.amp import custom_bwd, custom_fwd +from torch.autograd import Function + +amp_fwd = functools.partial(custom_fwd, device_type="cuda") +amp_bwd = functools.partial(custom_bwd, device_type="cuda") + +from .._libnatten import ( + reference_na1d_backward, + reference_na1d_forward, + reference_na2d_backward, + reference_na2d_forward, + reference_na3d_backward, + reference_na3d_forward, +) +from ..types import ( + CausalArg1DTypeOrDed, + CausalArg2DTypeOrDed, + CausalArg3DTypeOrDed, + CausalArgType, + CausalArgTypeOrDed, + Dimension1DTypeOrDed, + Dimension2DTypeOrDed, + Dimension3DTypeOrDed, + DimensionType, + DimensionTypeOrDed, + NoneType, +) +from ..utils import log +from ..utils.checks import ( + additional_kv_tensor_checks, + check_all_args, + check_args_against_input, + na_tensor_checks, +) + +logger = log.get_logger(__name__) + + +def make_reference_fna_autograd_fn(na_dim): + assert na_dim in [1, 2, 3] + + FORWARD_OPS = { + 1: reference_na1d_forward, + 2: reference_na2d_forward, + 3: reference_na3d_forward, + } + + BACKWARD_OPS = { + 1: reference_na1d_backward, + 2: reference_na2d_backward, + 3: reference_na3d_backward, + } + + class ReferenceFnaGenericAutogradFn(Function): + @staticmethod + @amp_fwd + def forward( + ctx, + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size: DimensionType, + stride: DimensionType, + dilation: DimensionType, + is_causal: CausalArgType, + scale: float, + qkv_shape: DimensionType, + num_extra_kv: int, + ) -> Tuple[Tensor, Tensor]: + kernel_size, stride, dilation, is_causal = check_all_args( + na_dim, kernel_size, stride, dilation, is_causal + ) + + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + + assert query.dim() == value.dim() == 4 + assert query.shape[0] == value.shape[0] + + output, logsumexp = FORWARD_OPS[na_dim]( + query, + key, + value, + kernel_size, + stride, + dilation, + is_causal, + scale, + qkv_shape, + num_extra_kv, + ) + + ctx.save_for_backward(query, key, value, logsumexp, output) + ctx.kernel_size = kernel_size + ctx.stride = stride + ctx.dilation = dilation + ctx.is_causal = is_causal + ctx.scale = scale + ctx.qkv_shape = qkv_shape + ctx.num_extra_kv = num_extra_kv + + return output, logsumexp + + @staticmethod + @amp_bwd + def backward(ctx, grad_out: Tensor, grad_lse: Tensor) -> Tuple[ + Tensor, + Tensor, + Tensor, + NoneType, + NoneType, + NoneType, + NoneType, + NoneType, + NoneType, + NoneType, + ]: + query, key, value, logsumexp, output = ctx.saved_tensors + d_output = grad_out.contiguous() + + d_query, d_key, d_value = BACKWARD_OPS[na_dim]( + query, + key, + value, + output, + d_output, + logsumexp, + ctx.kernel_size, + ctx.stride, + ctx.dilation, + ctx.is_causal, + ctx.scale, + ctx.qkv_shape, + ctx.num_extra_kv, + ) + + return d_query, d_key, d_value, None, None, None, None, None, None, None + + return ReferenceFnaGenericAutogradFn + + +ReferenceFna1DAutogradFn = make_reference_fna_autograd_fn(1) +ReferenceFna2DAutogradFn = make_reference_fna_autograd_fn(2) +ReferenceFna3DAutogradFn = make_reference_fna_autograd_fn(3) + + +ReferenceFnaAutogradFns = { + 1: ReferenceFna1DAutogradFn, + 2: ReferenceFna2DAutogradFn, + 3: ReferenceFna3DAutogradFn, +} + + +def reference_fna_generic( + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size: DimensionTypeOrDed, + stride: DimensionTypeOrDed = 1, + dilation: DimensionTypeOrDed = 1, + is_causal: Optional[CausalArgTypeOrDed] = False, + scale: Optional[float] = None, + additional_keys: Optional[Tensor] = None, + additional_values: Optional[Tensor] = None, + return_lse: bool = False, +) -> Union[Tensor, Tuple[Tensor, Tensor]]: + + na_tensor_checks( + query, key, value, must_match_head_dims=False, supports_gqa_mqa=True + ) + additional_kv_tensor_checks( + query, + key, + value, + additional_keys, + additional_values, + must_match_head_dims=False, + supports_gqa_mqa=True, + ) + + na_dim = query.dim() - 3 # batch, heads, head_dim + + assert na_dim in [1, 2, 3] + + kernel_size, stride, dilation, is_causal = check_all_args( + na_dim, kernel_size, stride, dilation, is_causal + ) + + check_args_against_input( + query, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + ) + + scale = scale or query.shape[-1] ** -0.5 + + qkv_shape = query.shape[1 : 1 + na_dim] + + query = query.flatten(1, na_dim) + key = key.flatten(1, na_dim) + value = value.flatten(1, na_dim) + + num_extra_kv = 0 + if additional_keys is not None and additional_values is not None: + num_extra_kv = additional_keys.shape[1] + key = torch.cat([key, additional_keys], dim=1) + value = torch.cat([value, additional_values], dim=1) + + output, lse = ReferenceFnaAutogradFns[na_dim].apply( + query, + key, + value, + kernel_size, + stride, + dilation, + is_causal, + scale, + qkv_shape, + num_extra_kv, + ) + output = output.reshape( + query.shape[0], *qkv_shape, query.shape[-2], value.shape[-1] + ) + lse = lse.reshape(query.shape[0], *qkv_shape, query.shape[-2]) + + if return_lse: + return output, lse + + return output + + +def na1d_reference( + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size: Dimension1DTypeOrDed, + stride: Dimension1DTypeOrDed = 1, + dilation: Dimension1DTypeOrDed = 1, + is_causal: Optional[CausalArg1DTypeOrDed] = False, + scale: Optional[float] = None, + additional_keys: Optional[Tensor] = None, + additional_values: Optional[Tensor] = None, + return_lse: bool = False, +) -> Union[Tensor, Tuple[Tensor, Tensor]]: + return reference_fna_generic( + query=query, + key=key, + value=value, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + scale=scale, + additional_keys=additional_keys, + additional_values=additional_values, + return_lse=return_lse, + ) + + +def na2d_reference( + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size: Dimension2DTypeOrDed, + stride: Dimension2DTypeOrDed = 1, + dilation: Dimension2DTypeOrDed = 1, + is_causal: Optional[CausalArg2DTypeOrDed] = False, + scale: Optional[float] = None, + additional_keys: Optional[Tensor] = None, + additional_values: Optional[Tensor] = None, + return_lse: bool = False, +) -> Union[Tensor, Tuple[Tensor, Tensor]]: + return reference_fna_generic( + query=query, + key=key, + value=value, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + scale=scale, + additional_keys=additional_keys, + additional_values=additional_values, + return_lse=return_lse, + ) + + +def na3d_reference( + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size: Dimension3DTypeOrDed, + stride: Dimension3DTypeOrDed = 1, + dilation: Dimension3DTypeOrDed = 1, + is_causal: Optional[CausalArg3DTypeOrDed] = False, + scale: Optional[float] = None, + additional_keys: Optional[Tensor] = None, + additional_values: Optional[Tensor] = None, + return_lse: bool = False, +) -> Union[Tensor, Tuple[Tensor, Tensor]]: + return reference_fna_generic( + query=query, + key=key, + value=value, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + scale=scale, + additional_keys=additional_keys, + additional_values=additional_values, + return_lse=return_lse, + ) diff --git a/natten/torch-ext/natten/context.py b/natten/torch-ext/natten/context.py new file mode 100644 index 00000000..1793ac98 --- /dev/null +++ b/natten/torch-ext/natten/context.py @@ -0,0 +1,231 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# +from enum import Enum + +import torch + +from .utils import log + +logger = log.get_logger(__name__) + + +class MemoryUsagePreference(Enum): + Default = 0 + Strict = 1 + Unrestricted = 2 + + +class NattenContext: + is_deterministic_mode_enabled: bool = False + is_kv_parallelism_enabled: bool = True + training_memory_preference: MemoryUsagePreference = MemoryUsagePreference.Default + flex_compile_allowed: bool = False + flex_compile_backprop_allowed: bool = False + + @staticmethod + def reset(): + NattenContext.is_deterministic_mode_enabled = False + NattenContext.is_kv_parallelism_enabled = True + NattenContext.training_memory_preference = MemoryUsagePreference.Default + NattenContext.flex_compile_allowed = False + NattenContext.flex_compile_backprop_allowed = False + + +def set_memory_usage_preference(pref: str = "default"): + """Sets memory usage preference for KV parallelism in `"cutlass-fna"` and `"cutlass-fmha"` + backends. + + Args: + pref: Choices are `"default"`, `"strict"`, and `"unrestricted"`. + """ + if pref == "default": + NattenContext.training_memory_preference = MemoryUsagePreference.Default + elif pref == "strict": + NattenContext.training_memory_preference = MemoryUsagePreference.Strict + elif pref == "unrestricted": + NattenContext.training_memory_preference = MemoryUsagePreference.Unrestricted + else: + raise ValueError( + "natten.set_memory_usage_preference allows only one of three settings: " + "`default`, `strict`, and `unrestricted`." + ) + + +def get_memory_usage_preference() -> MemoryUsagePreference: + return NattenContext.training_memory_preference + + +def is_memory_usage_default() -> bool: + """Returns whether memory usage preference for KV parallelism in `"cutlass-fna"` and + `"cutlass-fmha"` backends is the default setting. + """ + return get_memory_usage_preference() == MemoryUsagePreference.Default + + +def is_memory_usage_strict() -> bool: + """Returns whether memory usage preference for KV parallelism in `"cutlass-fna"` and + `"cutlass-fmha"` backends is the *restricted* setting. + """ + return get_memory_usage_preference() == MemoryUsagePreference.Strict + + +def is_memory_usage_unrestricted() -> bool: + """Returns whether memory usage preference for KV parallelism in `"cutlass-fna"` and + `"cutlass-fmha"` backends is the *unrestricted* setting. + """ + return get_memory_usage_preference() == MemoryUsagePreference.Unrestricted + + +def use_deterministic_algorithms(mode: bool = True): + NattenContext.is_deterministic_mode_enabled = mode + if mode: + logger.warning( + "You're enabling NATTEN's deterministic mode. This mode does not " + "support auto-tuning, or training with positional biases. " + "For more information please refer to https://github.com/SHI-Labs/NATTEN/tree/main/docs" + ) + + +def are_deterministic_algorithms_enabled() -> bool: + return NattenContext.is_deterministic_mode_enabled + + +def use_kv_parallelism_in_fused_na(mode: bool = True): + """Sets guards for using KV Parallelism in backpropagation in `"cutlass-fna"`/`"cutlass-fmha"` + backends. + + Warning: + Disabling KV parallelism can significantly slow down training, particularly in + small-batch/head and large-token problems. + + Args: + mode: If `True`, allows KV parallelism (default setting), and otherwise disables it. + """ + if not mode: + NattenContext.is_kv_parallelism_enabled = False + return + + if torch.are_deterministic_algorithms_enabled(): + logger.warning( + "Attempted to enable KV parallelism in FNA, which is non-deterministic, " + "but PyTorch's deterministic flag has been enabled. Ignoring..." + ) + return + + if are_deterministic_algorithms_enabled(): + raise RuntimeError( + "You enabled NATTEN's deterministic mode, but attempted to " + "enable KV parallelism, which results in non-determinism. " + ) + + NattenContext.is_kv_parallelism_enabled = True + + +def is_kv_parallelism_in_fused_na_enabled() -> bool: + """Returns whether KV parallelism in `"cutlass-fna"` and `"cutlass-fmha"` backends is enabled.""" + return NattenContext.is_kv_parallelism_enabled + + +def is_flex_compile_allowed() -> bool: + """Returns whether compilation is allowed in `"flex-fna"` and `"flex-fmha"` backends.""" + return NattenContext.flex_compile_allowed + + +def is_flex_compile_backprop_allowed() -> bool: + """Returns whether compilation for backpropagation is allowed in `"flex-fna"` and `"flex-fmha"` + backends. + """ + return NattenContext.flex_compile_backprop_allowed + + +def allow_flex_compile(mode: bool = True, backprop: bool = False): + """Sets guards for Flex Attention + `torch.compile`. + + Allows using our Flex FNA / Flex FMHA backends with `torch.compile`, meaning you can + pass `torch_compile=True` to the `na{1,2,3}d` or `attention` operation, along with + `backend="flex-fna"`/`backend="flex-fmha"`, and NATTEN will compile the block-sparse mask, as + well as the attention operation using `torch.compile` for you. + + Warning: + We have been *unable to verify the correctness* of this setting under all of our use + cases. We are working on raising this issue with PyTorch directly, but until then we strongly + recommend exercising caution when using this feature. + + Danger: backprop=True is strongly discouraged! + Allowing `torch.compile` for backpropagation (detected by checking + `tensor.requires_grad`) is guarded separately. We strongly recommend NOT using this setting, as + it can impact your training results. + + Args: + mode: If `True`, enable compilation for forward pass, otherwise disable. + backprop: If `True`, assuming compilation for forward pass is allowed, enable compilation + for backward pass, otherwise disable. + """ + if not mode: + NattenContext.flex_compile_allowed = False + NattenContext.flex_compile_backprop_allowed = False + + if not NattenContext.flex_compile_allowed: + logger.warning( + "You are enabling Flex Attention compilation in NATTEN. " + "NATTEN does not allow this by default, because we cannot verify Flex's correctness in all " + "scenarios through NATTEN's tests. By choosing to override this, you acknowledge that your " + "results may be affected significantly. If this was not intended, please call " + "natten.disable_flex_compile()" + "" + ) + + NattenContext.flex_compile_allowed = True + + if backprop: + if not NattenContext.flex_compile_backprop_allowed: + logger.warning( + "You are enabling using compiled Flex Attention to backpropagate. " + "NATTEN does not allow this by default, because we cannot verify Flex's correctness in all " + "scenarios through NATTEN's tests, and it is HIGHLY discouraged. By choosing to override " + "this, you acknowledge that your results may be heavily impacted significantly. " + "If this was not intended, please call " + "natten.disable_flex_compile_backprop()" + "" + ) + NattenContext.flex_compile_backprop_allowed = True + + +def allow_flex_compile_backprop(mode: bool = True): + """Sets guards for Flex Attention + `torch.compile` for backpropagation only. + + Args: + mode: If `True`, enable compilation for backprop (assuming forward compilation is already + enabled), otherwise disable. + """ + return allow_flex_compile(is_flex_compile_allowed(), mode) + + +def disable_flex_compile(): + """Disallow Flex Attention + `torch.compile` entirely.""" + return allow_flex_compile(False) + + +def disable_flex_compile_backprop(): + """Disallow Flex Attention + `torch.compile` for backpropagation entirely.""" + return allow_flex_compile(is_flex_compile_allowed(), False) diff --git a/natten/torch-ext/natten/functional.py b/natten/torch-ext/natten/functional.py new file mode 100644 index 00000000..d5cd1e42 --- /dev/null +++ b/natten/torch-ext/natten/functional.py @@ -0,0 +1,1151 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# +from typing import Dict, Optional, Tuple, Union + +import torch +from torch import Tensor + +from .attn_merge import merge_attentions +from .backends import ( + choose_backend, + choose_fmha_backend, + cutlass_blackwell_fmha, + cutlass_blackwell_fna_generic, + cutlass_fmha, + cutlass_fna_generic, + cutlass_hopper_fmha, + cutlass_hopper_fna_generic, + flex_fmha, + flex_fna_generic, +) +from .types import ( + CausalArg1DTypeOrDed, + CausalArg2DTypeOrDed, + CausalArg3DTypeOrDed, + CausalArgTypeOrDed, + Dimension1DType, + Dimension1DTypeOrDed, + Dimension2DType, + Dimension2DTypeOrDed, + Dimension3DType, + Dimension3DTypeOrDed, + DimensionType, + DimensionTypeOrDed, + KernelSchedule, +) +from .utils import log +from .utils.checks import ( + additional_kv_tensor_checks, + check_all_args, + check_args_against_input, + check_kernel_schedule, + fmha_tensor_checks, + is_self_attention, + na_tensor_checks, + varlen_tensor_checks, +) + +logger = log.get_logger(__name__) + + +# Standard Attention + + +def attention( + query: Tensor, + key: Tensor, + value: Tensor, + is_causal: bool = False, + scale: Optional[float] = None, + # varlen parameters + seqlens_Q: Optional[Tensor] = None, + seqlens_KV: Optional[Tensor] = None, + cumulative_seqlen_Q: Optional[Tensor] = None, + cumulative_seqlen_KV: Optional[Tensor] = None, + max_seqlen_Q: Optional[int] = None, + max_seqlen_KV: Optional[int] = None, + # backend parameters + backend: Optional[str] = None, + q_tile_size: Optional[int] = None, + kv_tile_size: Optional[int] = None, + backward_q_tile_size: Optional[int] = None, + backward_kv_tile_size: Optional[int] = None, + backward_kv_splits: Optional[int] = None, + backward_use_pt_reduction: bool = False, + run_persistent_kernel: bool = True, + kernel_schedule: Optional[Union[str, KernelSchedule]] = None, + torch_compile: bool = False, + return_lse: bool = False, +) -> Union[Tensor, Tuple[Tensor, Tensor]]: + """Runs standard dot product attention. + + This operation is used to implement neighborhood cross attention, in which we allow every + token to interact with some additional context (`additional_keys` and `additional_values` + tensors in [na1d][natten.na1d], [na2d][natten.na2d], and [na3d][natten.na3d]). + This operator is also used as a fast path for cases where neighborhood attention is equivalent + to self attention (not causal along any dims, and `kernel_size` is equal to the number of input + tokens). + + This operation does not call into PyTorch's SDPA, and only runs one of the NATTEN backends + (`cutlass-fmha`, `hopper-fmha`, `blackwell-fmha`, `flex-fmha`). Reasons for that include being + able to control performance-related arguments, return logsumexp, and more. + For more information refer to [backends](backends.md). + + Causal mask, and Variable length (varlen) Attention are also supported in some backends + (`cutlass-fmha` and `blackwell-fmha`). + + Varlen Attention is only supported for the sequence-packed layout: QKV tensors have batch size + 1, and tokens from different batches are concatenated without any padding along the sequence + dimension. Sequence lengths for different batches can be provided in two ways: + 1. `seqlens_Q` and `seqlens_KV` (less efficient): only provide the sequence lengths as + integer tensors (must be on the same device as QKV), and NATTEN will compute cumulative + and maximum sequence lengths on each call. + This is **incompatible** with full-graph `torch.compile` since it requires a + synchronization. + 2. `cumulative_seqlen_{Q,KV}` and `max_seqlen_{Q,KV}` (more efficient): + compute cumulative and maximum sequence lengths. `cumulative_seqlen_{Q,KV}` are integer + tensors on the same device as QKV containing the cumulative sum of `seqlens_{Q,KV}`, + with an additional `0` element in the beginning, therefore sized `batch+1`. + `max_seqlen_{Q,KV}` are integers (not Tensors) that represent the maximum sequence + lengths for Q and KV among all sequence batches. + You can use `natten.utils.varlen.generate_varlen_parameters` to generate these + parameters: + ```python3 + from .utils.varlen import generate_varlen_parameters + ( + cumulative_seqlen_Q, + cumulative_seqlen_KV, + max_seqlen_Q, + max_seqlen_KV, + ) = generate_varlen_parameters(q, k, v, seqlens_Q, seqlens_KV) + ``` + As long as `generate_varlen_parameters` is called ahead of torch.compiling the model, it + is supported without any graph breaks. + + GQA/MQA support (`heads != heads_kv`) is available. For now, `blackwell-fmha` and + `flex-fmha` support GQA/MQA natively, and `cutlass-fmha` and `hopper-fmha` support it with + explicit repeats (increases memory usage and runtime). + + Parameters: + query (Tensor): 4-D query tensor, with the heads last layout + (`[batch, seqlen, heads, head_dim]`) + + key (Tensor): 4-D key tensor, with the heads last layout + (`[batch, seqlen_kv, heads_kv, head_dim]`) + + value (Tensor): 4-D value tensor, with the heads last layout + (`[batch, seqlen_kv, heads_kv, head_dim_v]`) + + is_causal (bool): Toggle causal masking. Defaults to `False` (bi-directional). + + scale (float): Attention scale. Defaults to `head_dim ** -0.5`. + + seqlens_Q (Optional[Tensor]): (varlen) Optional 1-D tensor with size `batch` + indicating the number of query tokens in each batch. Must be passed together with + `seqlens_KV`. + + seqlens_KV (Optional[Tensor]): (varlen) Optional 1-D tensor with size `batch` + indicating the number of key/value tokens in each batch. Must be passed together with + `seqlens_Q`. + + cumulative_seqlen_Q (Optional[Tensor]): (varlen) Optional 1-D tensor with size `batch + 1` + indicating the cumulative sum of number of query tokens in each batch, with an + additional 0 element in the beginning. Must be passed together with + `cumulative_seqlen_KV` and `max_seqlen_{Q,KV}`. + + cumulative_seqlen_KV (Optional[Tensor]): (varlen) Optional 1-D tensor with size `batch + 1` + indicating the cumulative sum of number of key/value tokens in each batch, with an + additional 0 element in the beginning. Must be passed together with + `cumulative_seqlen_Q` and `max_seqlen_{Q,KV}`. + + max_seqlen_Q (Optional[int]): (varlen) Optional integer indicating the maximum query + sequence length in all batches. Must be passed together with `cumulative_seqlen_{Q,KV}` + and `max_seqlen_KV`. + + max_seqlen_KV (Optional[int]): (varlen) Optional integer indicating the maximum key/value + sequence length in all batches. Must be passed together with `cumulative_seqlen_{Q,KV}` + and `max_seqlen_Q`. + + Other Parameters: + backend (str): Backend implementation to run with. Choices are: `None` (pick the best + available one), `"cutlass-fmha"`, `"hopper-fmha"`, `"blackwell-fmha"`, `"flex-fmha"`. + Refer to [backends](backends.md) for more information. + + q_tile_size (int): Tile size along query sequence length in the forward pass kernel. + You can use [profiler](profiler.md) to find valid choices for your use case. + + kv_tile_size (int): Tile size along key/value sequence length in the forward pass kernel. + You can use [profiler](profiler.md) to find valid choices for your use case. + + backward_q_tile_size (int): Tile size along query sequence length in the backward pass + kernel. This is ignored by `"flex-fmha"`. + You can use [profiler](profiler.md) to find valid choices for your use case. + + backward_kv_tile_size (int): Tile size along key/value sequence length in the backward pass + kernel. This is ignored by `"flex-fmha"`. + You can use [profiler](profiler.md) to find valid choices for your use case. + + backward_kv_splits (int): Number of key/value tiles allowed to work in parallel in the + backward pass kernel. This is only respected by the `"cutlass-fmha"` backend, only when + [KV parallelism](context.md#kv-parallelism-in-fna) is enabled. + + backward_use_pt_reduction (bool): Whether to use PyTorch eager for computing the `dO * O` + product required by the backward pass, over the CUTLASS kernel. This only applies to + the `"cutlass-fmha"` backend. + + run_persistent_kernel (bool): Whether to use persistent tile scheduling in the forward pass + kernel. This only applies to the `"blackwell-fmha"` backend. + + kernel_schedule (Optional[str]): Kernel type (Hopper architecture only). Choices are + `None`: pick the default, `"non"` (non-persistent), `"coop"` (warp-specialized + cooperative), or `"pp"` (warp-specialized ping-ponging). Refer to + [Hopper FMHA/FNA backend](backends.md#hopper-fna-fmha) for more information. + + torch_compile (bool): Applies only to the `"flex-fmha"` backend. Whether or not to JIT + compile the attention kernel. Due to this being an experimental feature in PyTorch, we + do not recommend it, and it is guarded by context flags. Read more in + [Flex Attention + `torch.compile`](context.md#flex-attention-torchcompile). + + return_lse (bool): Whether or not to return the `logsumexp` tensor. `logsumexp` can be used + in the backward pass, and for [attention merging][natten.merge_attentions]. + + Returns: + output (Tensor): 4-D output tensor, with the heads last layout + (`[batch, seqlen, heads, head_dim_v]`). + + logsumexp (Tensor): only returned when `return_lse=True`. 3-D logsumexp tensor, with the + heads last layout (`[batch, seqlen, heads]`). + """ + + fmha_tensor_checks(query, key, value) + + ( + cumulative_seqlen_Q, + cumulative_seqlen_KV, + max_seqlen_Q, + max_seqlen_KV, + ) = varlen_tensor_checks( + query=query, + key=key, + value=value, + seqlens_Q=seqlens_Q, + seqlens_KV=seqlens_KV, + cumulative_seqlen_Q=cumulative_seqlen_Q, + cumulative_seqlen_KV=cumulative_seqlen_KV, + max_seqlen_Q=max_seqlen_Q, + max_seqlen_KV=max_seqlen_KV, + ) + is_varlen = cumulative_seqlen_Q is not None + + scale = scale or query.shape[-1] ** -0.5 + + kernel_schedule = check_kernel_schedule(kernel_schedule) + + backend = backend or choose_fmha_backend( + query, + key, + value, + is_causal=is_causal, + is_varlen=is_varlen, + torch_compile=torch_compile, + ) + + if backend == "blackwell-fmha": + return cutlass_blackwell_fmha( + query=query, + key=key, + value=value, + is_causal=is_causal, + scale=scale, + cumulative_seqlen_Q=cumulative_seqlen_Q, + cumulative_seqlen_KV=cumulative_seqlen_KV, + max_seqlen_Q=max_seqlen_Q, + max_seqlen_KV=max_seqlen_KV, + q_tile_size=q_tile_size, + kv_tile_size=kv_tile_size, + backward_q_tile_size=backward_q_tile_size, + backward_kv_tile_size=backward_kv_tile_size, + run_persistent_kernel=run_persistent_kernel, + return_lse=return_lse, + ) + + if backend == "hopper-fmha": + return cutlass_hopper_fmha( + query=query, + key=key, + value=value, + is_causal=is_causal, + scale=scale, + q_tile_size=q_tile_size, + kv_tile_size=kv_tile_size, + backward_q_tile_size=backward_q_tile_size, + backward_kv_tile_size=backward_kv_tile_size, + kernel_schedule=kernel_schedule, + return_lse=return_lse, + cumulative_seqlen_Q=cumulative_seqlen_Q, + cumulative_seqlen_KV=cumulative_seqlen_KV, + max_seqlen_Q=max_seqlen_Q, + max_seqlen_KV=max_seqlen_KV, + ) + + elif backend == "cutlass-fmha": + return cutlass_fmha( + query=query, + key=key, + value=value, + is_causal=is_causal, + scale=scale, + q_tile_size=q_tile_size, + kv_tile_size=kv_tile_size, + backward_q_tile_size=backward_q_tile_size, + backward_kv_tile_size=backward_kv_tile_size, + backward_kv_splits=backward_kv_splits, + backward_use_pt_reduction=backward_use_pt_reduction, + return_lse=return_lse, + cumulative_seqlen_Q=cumulative_seqlen_Q, + cumulative_seqlen_KV=cumulative_seqlen_KV, + max_seqlen_Q=max_seqlen_Q, + max_seqlen_KV=max_seqlen_KV, + ) + + elif backend == "flex-fmha": + return flex_fmha( + query=query, + key=key, + value=value, + is_causal=is_causal, + scale=scale, + q_tile_size=q_tile_size, + kv_tile_size=kv_tile_size, + torch_compile=torch_compile, + return_lse=return_lse, + cumulative_seqlen_Q=cumulative_seqlen_Q, + cumulative_seqlen_KV=cumulative_seqlen_KV, + max_seqlen_Q=max_seqlen_Q, + max_seqlen_KV=max_seqlen_KV, + ) + + raise NotImplementedError(f"Unrecognized NATTEN FMHA backend {backend}.") + + +# Neighborhood Attention + + +def neighborhood_attention_generic( + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size: DimensionTypeOrDed, + stride: DimensionTypeOrDed = 1, + dilation: DimensionTypeOrDed = 1, + is_causal: Optional[CausalArgTypeOrDed] = False, + scale: Optional[float] = None, + additional_keys: Optional[Tensor] = None, + additional_values: Optional[Tensor] = None, + attention_kwargs: Optional[Dict] = None, + # Perf-related args + backend: Optional[str] = None, + q_tile_shape: Optional[DimensionType] = None, + kv_tile_shape: Optional[DimensionType] = None, + backward_q_tile_shape: Optional[DimensionType] = None, + backward_kv_tile_shape: Optional[DimensionType] = None, + backward_kv_splits: Optional[DimensionType] = None, + backward_use_pt_reduction: bool = False, + run_persistent_kernel: bool = True, + kernel_schedule: Optional[Union[str, KernelSchedule]] = None, + torch_compile: bool = False, + return_lse: bool = False, +) -> Union[Tensor, Tuple[Tensor, Tensor]]: + + na_tensor_checks(query, key, value) + additional_kv_tensor_checks(query, key, value, additional_keys, additional_values) + kernel_schedule = check_kernel_schedule(kernel_schedule) + + na_dim = query.dim() - 3 # batch, heads, head_dim + + assert na_dim in [1, 2, 3] + + kernel_size, stride, dilation, is_causal = check_all_args( + na_dim, kernel_size, stride, dilation, is_causal + ) + + check_args_against_input( + query, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + ) + + has_additional_attention = ( + additional_keys is not None and additional_values is not None + ) + + if is_self_attention( + query, + kernel_size=kernel_size, + is_causal=is_causal, + has_additional_attention=has_additional_attention, + ): + logger.debug( + f"{query.shape=} with {kernel_size=}, {has_additional_attention=} and {is_causal=} is " + "self attention. Calling attention instead of neighborhood attention directly." + ) + + query_shape = query.shape + query = query.flatten(1, na_dim) + key = key.flatten(1, na_dim) + value = value.flatten(1, na_dim) + + if has_additional_attention: + assert additional_keys is not None + assert additional_values is not None + key = torch.cat([key, additional_keys], dim=1) + value = torch.cat([value, additional_values], dim=1) + + attn_kwargs = attention_kwargs or {} + out, lse = attention( + query, + key, + value, + is_causal=is_causal[0], # NOTE: special case + scale=scale, + return_lse=True, + **attn_kwargs, + ) + lse_shape = [s for s in query_shape[:-1]] + output_shape = lse_shape + [value.shape[-1]] + out = out.reshape(*output_shape) + lse = lse.reshape(*lse_shape) + + if return_lse: + return out, lse + + return out + + scale = scale or query.shape[-1] ** -0.5 + + backend = backend or choose_backend(query, key, value, torch_compile=torch_compile) + + if backend == "blackwell-fna": + output, lse = cutlass_blackwell_fna_generic( + query=query, + key=key, + value=value, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + scale=scale, + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + backward_q_tile_shape=backward_q_tile_shape, + backward_kv_tile_shape=backward_kv_tile_shape, + run_persistent_kernel=run_persistent_kernel, + return_lse=True, + ) + + elif backend == "hopper-fna": + output, lse = cutlass_hopper_fna_generic( + query=query, + key=key, + value=value, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + scale=scale, + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + backward_q_tile_shape=backward_q_tile_shape, + backward_kv_tile_shape=backward_kv_tile_shape, + kernel_schedule=kernel_schedule, + return_lse=True, + ) + + elif backend == "cutlass-fna": + output, lse = cutlass_fna_generic( + query=query, + key=key, + value=value, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + scale=scale, + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + backward_q_tile_shape=backward_q_tile_shape, + backward_kv_tile_shape=backward_kv_tile_shape, + backward_kv_splits=backward_kv_splits, + backward_use_pt_reduction=backward_use_pt_reduction, + return_lse=True, + ) + + elif backend == "flex-fna": + output, lse = flex_fna_generic( + query=query, + key=key, + value=value, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + scale=scale, + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + torch_compile=torch_compile, + return_lse=True, + ) + + else: + raise NotImplementedError(f"Unrecognized NATTEN backend {backend}.") + + if has_additional_attention: + assert additional_keys is not None + assert additional_values is not None + + attention_kwargs = attention_kwargs or {} + if "torch_compile" not in attention_kwargs: + attention_kwargs["torch_compile"] = torch_compile + + additional_output, additional_lse = attention( + query.flatten(1, na_dim), + additional_keys, + additional_values, + is_causal=False, + scale=scale, + return_lse=True, + **attention_kwargs, + ) + + # NOTE: Flex unfused should not use the autograd fix + is_flex = backend == "flex-fna" or ( + "backend" in attention_kwargs and attention_kwargs["backend"] == "flex-fmha" + ) + use_autograd_fix = not is_flex or torch_compile + + merged_output, merged_lse = merge_attentions( + [output.flatten(1, na_dim), additional_output], + [lse.flatten(1, na_dim), additional_lse], + use_autograd_fix=use_autograd_fix, + ) + merged_output = merged_output.reshape(output.shape) + merged_lse = merged_lse.reshape(output.shape[:-1]) + + if return_lse: + return merged_output, merged_lse + + return merged_output + + if return_lse: + return output, lse + + return output + + +def na1d( + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size: Dimension1DTypeOrDed, + stride: Dimension1DTypeOrDed = 1, + dilation: Dimension1DTypeOrDed = 1, + is_causal: Optional[CausalArg1DTypeOrDed] = False, + scale: Optional[float] = None, + additional_keys: Optional[Tensor] = None, + additional_values: Optional[Tensor] = None, + attention_kwargs: Optional[Dict] = None, + backend: Optional[str] = None, + q_tile_shape: Optional[Dimension1DType] = None, + kv_tile_shape: Optional[Dimension1DType] = None, + backward_q_tile_shape: Optional[Dimension1DType] = None, + backward_kv_tile_shape: Optional[Dimension1DType] = None, + backward_kv_splits: Optional[Dimension1DType] = None, + backward_use_pt_reduction: bool = False, + run_persistent_kernel: bool = True, + kernel_schedule: Optional[Union[str, KernelSchedule]] = None, + torch_compile: bool = False, + return_lse: bool = False, +) -> Union[Tensor, Tuple[Tensor, Tensor]]: + """Computes 1-D neighborhood attention. + + GQA/MQA support (`heads != heads_kv`) is available. For now, `blackwell-fna` and + `flex-fna` support GQA/MQA natively, and `cutlass-fna` and `hopper-fna` support it with + explicit repeats (increases memory usage and runtime). + + Parameters: + query (Tensor): 4-D query tensor, with the heads last layout + (`[batch, seqlen, heads, head_dim]`) + + key (Tensor): 4-D key tensor, with the heads last layout + (`[batch, seqlen, heads_kv, head_dim]`) + + value (Tensor): 4-D value tensor, with the heads last layout + (`[batch, seqlen, heads_kv, head_dim_v]`) + + kernel_size (Tuple[int] | int): Neighborhood window (kernel) size. + + !!! note + `kernel_size` must be smaller than or equal to `seqlen`. + + stride (Tuple[int] | int): Sliding window step size. Defaults to `1` (standard sliding + window). + + !!! note + `stride` must be smaller than or equal to `kernel_size`. + When `stride == kernel_size`, there will be no overlap between sliding windows, + which is equivalent to blocked attention (a.k.a. + [window self attention](https://arxiv.org/abs/2103.14030)). + + dilation (Tuple[int] | int): Dilation step size. Defaults to `1` (standard sliding window). + + !!! note + The product of `dilation` and `kernel_size` must be smaller than or equal to + `seqlen`. + + is_causal (Tuple[bool] | bool): Toggle causal masking. Defaults to `False` + (bi-directional). + + scale (float): Attention scale. Defaults to `head_dim ** -0.5`. + + additional_keys: `None` or 4-D key tensor, with the heads last layout + (`[batch, seqlen_kv, heads_kv, head_dim]`), corresponding to key tokens from some + additional context. Used when performing neighborhood cross-attention, where `query` + tokens attend to their neighborhood, as well as some fixed additional set of tokens. + + additional_values: `None` or 4-D value tensor, with the heads last layout + (`[batch, seqlen_kv, heads_kv, head_dim_v]`), corresponding to value tokens from some + additional context. Used when performing neighborhood cross-attention, where `query` + tokens attend to their neighborhood, as well as some fixed additional set of tokens. + + !!! note + `additional_keys` and `additional_values` must both either be `Tensor`s, or both + `None`s, and must match in shape. + + Other Parameters: + backend (str): Backend implementation to run with. Choices are: `None` (pick the best + available one), `"cutlass-fna"`, `"hopper-fna"`, `"blackwell-fna"`, `"flex-fna"`. + Refer to [backends](backends.md) for more information. + + q_tile_shape (Tuple[int]): 1-D Tile shape for the query token layout in the forward pass + kernel. You can use [profiler](profiler.md) to find valid choices for your use case, + and search for the best combination. + + kv_tile_shape (Tuple[int]): 1-D Tile shape for the key-value token layout in the forward + pass kernel. You can use [profiler](profiler.md) to find valid choices for your use + case, and search for the best combination. + + backward_q_tile_shape (Tuple[int]): 1-D Tile shape for the query token layout in the + backward pass kernel. This is ignored by `"flex-fna"`. + You can use [profiler](profiler.md) to find valid choices for your use case, and + search for the best combination. + + backward_kv_tile_shape (Tuple[int]): 1-D Tile shape for the key/value token layout in the + backward pass kernel. This is ignored by `"flex-fna"`. + You can use [profiler](profiler.md) to find valid choices for your use case, and + search for the best combination. + + backward_kv_splits (Tuple[int]): Number of key/value tiles allowed to work in parallel in + the backward pass kernel. Like tile shapes, this is a tuple and not an integer for + neighborhood attention operations, and the size of the tuple corresponds to the number + of dimensions / rank of the layout of tokens. This is only respected by the + `"cutlass-fna"` backend, and only when + [KV parallelism](context.md#kv-parallelism-in-fna) is enabled. + + backward_use_pt_reduction (bool): Whether to use PyTorch eager for computing the `dO * O` + product required by the backward pass, over the CUTLASS kernel. This only applies to + the `"cutlass-fna"` backend. + + run_persistent_kernel (bool): Whether to use persistent tile scheduling in the forward pass + kernel. This only applies to the `"blackwell-fna"` backend. + + kernel_schedule (Optional[str]): Kernel type (Hopper architecture only). Choices are + `None`: pick the default, `"non"` (non-persistent), `"coop"` (warp-specialized + cooperative), or `"pp"` (warp-specialized ping-ponging). Refer to + [Hopper FMHA/FNA backend](backends.md#hopper-fna-fmha) for more information. + + torch_compile (bool): Applies only to the `"flex-fna"` backend. Whether or not to JIT + compile the attention kernel. Due to this being an experimental feature in PyTorch, we + do not recommend it, and it is guarded by context flags. Read more in + [Flex Attention + `torch.compile`](context.md#flex-attention-torchcompile). + + attention_kwargs: arguments to the [attention][natten.attention] operator, if used to + implement neighborhood cross-attention, or self attention as a fast path for + neighborhood attention. + + If `additional_{keys,values}` are specified, NATTEN usually performs a separate + cross-attention using our [attention][natten.attention] operator, and + [merges][natten.merge_attentions] the results. + + If for a given use case, the neighborhood attention problem is equivalent to self + attention (not causal, `kernel_size == seqlen`), NATTEN will also attempt to directly + use [attention][natten.attention]. + + You can override arguments to [attention][natten.attention] by passing a + dictionary here. + + !!! example + ```python + out = na1d( + q, k, v, kernel_size=kernel_size, + ..., + attention_kwargs={ + "backend": "blackwell-fmha", + "run_persistent_kernel": True, + } + ) + ``` + + return_lse (bool): Whether or not to return the `logsumexp` tensor. `logsumexp` can be used + in the backward pass, and for [attention merging][natten.merge_attentions]. + + Returns: + output (Tensor): 4-D output tensor, with the heads last layout + (`[batch, seqlen, heads, head_dim_v]`). + + logsumexp (Tensor): only returned when `return_lse=True`. 3-D logsumexp tensor, with the + heads last layout (`[batch, seqlen, heads]`). + """ + return neighborhood_attention_generic( + query=query, + key=key, + value=value, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + scale=scale, + additional_keys=additional_keys, + additional_values=additional_values, + attention_kwargs=attention_kwargs, + backend=backend, + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + backward_q_tile_shape=backward_q_tile_shape, + backward_kv_tile_shape=backward_kv_tile_shape, + backward_kv_splits=backward_kv_splits, + backward_use_pt_reduction=backward_use_pt_reduction, + run_persistent_kernel=run_persistent_kernel, + kernel_schedule=kernel_schedule, + torch_compile=torch_compile, + return_lse=return_lse, + ) + + +def na2d( + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size: Dimension2DTypeOrDed, + stride: Dimension2DTypeOrDed = 1, + dilation: Dimension2DTypeOrDed = 1, + is_causal: Optional[CausalArg2DTypeOrDed] = False, + scale: Optional[float] = None, + additional_keys: Optional[Tensor] = None, + additional_values: Optional[Tensor] = None, + attention_kwargs: Optional[Dict] = None, + backend: Optional[str] = None, + q_tile_shape: Optional[Dimension2DType] = None, + kv_tile_shape: Optional[Dimension2DType] = None, + backward_q_tile_shape: Optional[Dimension2DType] = None, + backward_kv_tile_shape: Optional[Dimension2DType] = None, + backward_kv_splits: Optional[Dimension2DType] = None, + backward_use_pt_reduction: bool = False, + run_persistent_kernel: bool = True, + kernel_schedule: Optional[Union[str, KernelSchedule]] = None, + torch_compile: bool = False, + return_lse: bool = False, +) -> Union[Tensor, Tuple[Tensor, Tensor]]: + """Computes 2-D neighborhood attention. + + GQA/MQA support (`heads != heads_kv`) is available. For now, `blackwell-fna` and + `flex-fna` support GQA/MQA natively, and `cutlass-fna` and `hopper-fna` support it with + explicit repeats (increases memory usage and runtime). + + Parameters: + query (Tensor): 2-D query tensor, with the heads last layout: + `[batch, X, Y, heads, head_dim]`, where token layout shape (feature map shape) is + `(X, Y)`. + + key (Tensor): 2-D key tensor, with the heads last layout: + `[batch, X, Y, heads_kv, head_dim]`, where token layout shape (feature map shape) is + `(X, Y)`. + + value (Tensor): 2-D value tensor, with the heads last layout: + `[batch, X, Y, heads_kv, head_dim_v]`, where token layout shape (feature map shape) is + `(X, Y)`. + + kernel_size (Tuple[int, int] | int): Neighborhood window (kernel) size/shape. If an + integer, it will be repeated for all 2 dimensions. For example `kernel_size=3` is + reinterpreted as `kernel_size=(3, 3)`. + + !!! note + `kernel_size` must be smaller than or equal to token layout shape (`(X, Y)`) along + every dimension. + + stride (Tuple[int, int] | int): Sliding window step size/shape. Defaults to `1` (standard + sliding window). If an integer, it will be repeated for all 2 dimensions. For example + `stride=2` is reinterpreted as `stride=(2, 2)`. + + !!! note + `stride` must be smaller than or equal to `kernel_size` along every dimension. + When `stride == kernel_size`, there will be no overlap between sliding windows, + which is equivalent to blocked attention (a.k.a. + [window self attention](https://arxiv.org/abs/2103.14030)). + + dilation (Tuple[int, int] | int): Dilation step size/shape. Defaults to `1` (standard + sliding window). If an integer, it will be repeated for all 2 dimensions. For example + `dilation=4` is reinterpreted as `dilation=(4, 4)`. + + !!! note + The product of `dilation` and `kernel_size` must be smaller than or equal to + token layout shape (`(X, Y)`) along every dimension. + + is_causal (Tuple[bool, bool] | bool): Toggle causal masking. Defaults to `False` + (bi-directional). If a boolean, it will be repeated for all 2 dimensions. For example + `is_causal=True` is reinterpreted as `is_causal=(True, True)`. + + scale (float): Attention scale. Defaults to `head_dim ** -0.5`. + + additional_keys: `None` or 4-D key tensor, with the heads last layout + (`[batch, seqlen_kv, heads_kv, head_dim]`), corresponding to key tokens from some + additional context. Used when performing neighborhood cross-attention, where `query` + tokens attend to their neighborhood, as well as some fixed additional set of tokens. + + additional_values: `None` or 4-D value tensor, with the heads last layout + (`[batch, seqlen_kv, heads_kv, head_dim_v]`), corresponding to value tokens from some + additional context. Used when performing neighborhood cross-attention, where `query` + tokens attend to their neighborhood, as well as some fixed additional set of tokens. + + !!! note + `additional_keys` and `additional_values` must both either be `Tensor`s, or both + `None`s, and must match in shape. + + Other Parameters: + backend (str): Backend implementation to run with. Choices are: `None` (pick the best + available one), `"cutlass-fna"`, `"hopper-fna"`, `"blackwell-fna"`, `"flex-fna"`. + Refer to [backends](backends.md) for more information. + + q_tile_shape (Tuple[int, int]): 2-D Tile shape for the query token layout in the forward + pass kernel. You can use [profiler](profiler.md) to find valid choices for your use + case, and search for the best combination. + + kv_tile_shape (Tuple[int, int]): 2-D Tile shape for the key-value token layout in the + forward pass kernel. You can use [profiler](profiler.md) to find valid choices for your + use case, and search for the best combination. + + backward_q_tile_shape (Tuple[int, int]): 2-D Tile shape for the query token layout in the + backward pass kernel. This is ignored by `"flex-fna"`. + You can use [profiler](profiler.md) to find valid choices for your use case, and + search for the best combination. + + backward_kv_tile_shape (Tuple[int, int]): 2-D Tile shape for the key/value token layout in + the backward pass kernel. This is ignored by `"flex-fna"`. + You can use [profiler](profiler.md) to find valid choices for your use case, and + search for the best combination. + + backward_kv_splits (Tuple[int, int]): Number of key/value tiles allowed to work in parallel + in the backward pass kernel. Like tile shapes, this is a tuple and not an integer for + neighborhood attention operations, and the size of the tuple corresponds to the number + of dimensions / rank of the layout of tokens. This is only respected by the + `"cutlass-fna"` backend, and only when + [KV parallelism](context.md#kv-parallelism-in-fna) is enabled. + + backward_use_pt_reduction (bool): Whether to use PyTorch eager for computing the `dO * O` + product required by the backward pass, over the CUTLASS kernel. This only applies to + the `"cutlass-fna"` backend. + + run_persistent_kernel (bool): Whether to use persistent tile scheduling in the forward pass + kernel. This only applies to the `"blackwell-fna"` backend. + + kernel_schedule (Optional[str]): Kernel type (Hopper architecture only). Choices are + `None`: pick the default, `"non"` (non-persistent), `"coop"` (warp-specialized + cooperative), or `"pp"` (warp-specialized ping-ponging). Refer to + [Hopper FMHA/FNA backend](backends.md#hopper-fna-fmha) for more information. + + torch_compile (bool): Applies only to the `"flex-fna"` backend. Whether or not to JIT + compile the attention kernel. Due to this being an experimental feature in PyTorch, we + do not recommend it, and it is guarded by context flags. Read more in + [Flex Attention + `torch.compile`](context.md#flex-attention-torchcompile). + + attention_kwargs: arguments to the [attention][natten.attention] operator, if used to + implement neighborhood cross-attention, or self attention as a fast path for + neighborhood attention. + + If `additional_{keys,values}` are specified, NATTEN usually performs a separate + cross-attention using our [attention][natten.attention] operator, and + [merges][natten.merge_attentions] the results. + + If for a given use case, the neighborhood attention problem is equivalent to self + attention (not causal along any dims, `kernel_size == (X, Y)`), NATTEN will also + attempt to directly use [attention][natten.attention]. + + You can override arguments to [attention][natten.attention] by passing a + dictionary here. + + !!! example + ```python + out = na2d( + q, k, v, kernel_size=kernel_size, + ..., + attention_kwargs={ + "backend": "blackwell-fmha", + "run_persistent_kernel": True, + } + ) + ``` + + return_lse (bool): Whether or not to return the `logsumexp` tensor. `logsumexp` can be used + in the backward pass, and for [attention merging][natten.merge_attentions]. + + Returns: + output (Tensor): 5-D output tensor, with the heads last layout + (`[batch, X, Y, heads, head_dim_v]`). + + logsumexp (Tensor): only returned when `return_lse=True`. 4-D logsumexp tensor, with the + heads last layout (`[batch, X, Y, heads]`). + """ + return neighborhood_attention_generic( + query=query, + key=key, + value=value, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + scale=scale, + additional_keys=additional_keys, + additional_values=additional_values, + attention_kwargs=attention_kwargs, + backend=backend, + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + backward_q_tile_shape=backward_q_tile_shape, + backward_kv_tile_shape=backward_kv_tile_shape, + backward_kv_splits=backward_kv_splits, + backward_use_pt_reduction=backward_use_pt_reduction, + run_persistent_kernel=run_persistent_kernel, + kernel_schedule=kernel_schedule, + torch_compile=torch_compile, + return_lse=return_lse, + ) + + +def na3d( + query: Tensor, + key: Tensor, + value: Tensor, + kernel_size: Dimension3DTypeOrDed, + stride: Dimension3DTypeOrDed = 1, + dilation: Dimension3DTypeOrDed = 1, + is_causal: Optional[CausalArg3DTypeOrDed] = False, + scale: Optional[float] = None, + additional_keys: Optional[Tensor] = None, + additional_values: Optional[Tensor] = None, + attention_kwargs: Optional[Dict] = None, + backend: Optional[str] = None, + q_tile_shape: Optional[Dimension3DType] = None, + kv_tile_shape: Optional[Dimension3DType] = None, + backward_q_tile_shape: Optional[Dimension3DType] = None, + backward_kv_tile_shape: Optional[Dimension3DType] = None, + backward_kv_splits: Optional[Dimension3DType] = None, + backward_use_pt_reduction: bool = False, + run_persistent_kernel: bool = True, + kernel_schedule: Optional[Union[str, KernelSchedule]] = None, + torch_compile: bool = False, + return_lse: bool = False, +) -> Union[Tensor, Tuple[Tensor, Tensor]]: + """Computes 3-D neighborhood attention. + + GQA/MQA support (`heads != heads_kv`) is available. For now, `blackwell-fna` and + `flex-fna` support GQA/MQA natively, and `cutlass-fna` and `hopper-fna` support it with + explicit repeats (increases memory usage and runtime). + + Parameters: + query (Tensor): 3-D query tensor, with the heads last layout: + `[batch, X, Y, Z, heads, head_dim]`, where token layout shape (feature map shape) is + `(X, Y, Z)`. + + key (Tensor): 3-D key tensor, with the heads last layout: + `[batch, X, Y, Z, heads_kv, head_dim]`, where token layout shape (feature map shape) is + `(X, Y, Z)`. + + value (Tensor): 3-D value tensor, with the heads last layout: + `[batch, X, Y, Z, heads_kv, head_dim_V]`, where token layout shape (feature map shape) is + `(X, Y, Z)`. + + kernel_size (Tuple[int, int, int] | int): Neighborhood window (kernel) size/shape. If an + integer, it will be repeated for all 3 dimensions. For example `kernel_size=3` is + reinterpreted as `kernel_size=(3, 3, 3)`. + + !!! note + `kernel_size` must be smaller than or equal to token layout shape (`(X, Y, Z)`) + along every dimension. + + stride (Tuple[int, int, int] | int): Sliding window step size/shape. Defaults to `1` + (standard sliding window). If an integer, it will be repeated for all 3 dimensions. + For example `stride=2` is reinterpreted as `stride=(2, 2, 2)`. + + !!! note + `stride` must be smaller than or equal to `kernel_size` along every dimension. + When `stride == kernel_size`, there will be no overlap between sliding windows, + which is equivalent to blocked attention (a.k.a. + [window self attention](https://arxiv.org/abs/2103.14030)). + + dilation (Tuple[int, int, int] | int): Dilation step size/shape. Defaults to `1` (standard + sliding window). If an integer, it will be repeated for all 3 dimensions. For example + `dilation=4` is reinterpreted as `dilation=(4, 4, 4)`. + + !!! note + The product of `dilation` and `kernel_size` must be smaller than or equal to + token layout shape (`(X, Y, Z)`) along every dimension. + + is_causal (Tuple[bool, bool, bool] | bool): Toggle causal masking. Defaults to `False` + (bi-directional). If a boolean, it will be repeated for all 3 dimensions. For example + `is_causal=True` is reinterpreted as `is_causal=(True, True, True)`. + + scale (float): Attention scale. Defaults to `head_dim ** -0.5`. + + additional_keys: `None` or 4-D key tensor, with the heads last layout + (`[batch, seqlen_kv, heads_kv, head_dim_v]`), corresponding to key tokens from some + additional context. Used when performing neighborhood cross-attention, where `query` + tokens attend to their neighborhood, as well as some fixed additional set of tokens. + + additional_values: `None` or 4-D value tensor, with the heads last layout + (`[batch, seqlen_kv, heads_kv, head_dim_v]`), corresponding to value tokens from some + additional context. Used when performing neighborhood cross-attention, where `query` + tokens attend to their neighborhood, as well as some fixed additional set of tokens. + + !!! note + `additional_keys` and `additional_values` must both either be `Tensor`s, or both + `None`s, and must match in shape. + + Other Parameters: + backend (str): Backend implementation to run with. Choices are: `None` (pick the best + available one), `"cutlass-fna"`, `"hopper-fna"`, `"blackwell-fna"`, `"flex-fna"`. + Refer to [backends](backends.md) for more information. + + q_tile_shape (Tuple[int, int, int]): 3-D Tile shape for the query token layout in the + forward pass kernel. You can use [profiler](profiler.md) to find valid choices for your + use case, and search for the best combination. + + kv_tile_shape (Tuple[int, int, int]): 3-D Tile shape for the key-value token layout in the + forward pass kernel. You can use [profiler](profiler.md) to find valid choices for your + use case, and search for the best combination. + + backward_q_tile_shape (Tuple[int, int, int]): 3-D Tile shape for the query token layout in + the backward pass kernel. This is ignored by `"flex-fna"`. + You can use [profiler](profiler.md) to find valid choices for your use case, and + search for the best combination. + + backward_kv_tile_shape (Tuple[int, int, int]): 3-D Tile shape for the key/value token + layout in the backward pass kernel. This is ignored by `"flex-fna"`. + You can use [profiler](profiler.md) to find valid choices for your use case, + and search for the best combination. + + backward_kv_splits (Tuple[int, int, int]): Number of key/value tiles allowed to work in + parallel in the backward pass kernel. Like tile shapes, this is a tuple and not an + integer for neighborhood attention operations, and the size of the tuple corresponds to + the number of dimensions / rank of the layout of tokens. This is only respected by the + `"cutlass-fna"` backend, and only when + [KV parallelism](context.md#kv-parallelism-in-fna) is enabled. + + backward_use_pt_reduction (bool): Whether to use PyTorch eager for computing the `dO * O` + product required by the backward pass, over the CUTLASS kernel. This only applies to + the `"cutlass-fna"` backend. + + run_persistent_kernel (bool): Whether to use persistent tile scheduling in the forward pass + kernel. This only applies to the `"blackwell-fna"` backend. + + kernel_schedule (Optional[str]): Kernel type (Hopper architecture only). Choices are + `None`: pick the default, `"non"` (non-persistent), `"coop"` (warp-specialized + cooperative), or `"pp"` (warp-specialized ping-ponging). Refer to + [Hopper FMHA/FNA backend](backends.md#hopper-fna-fmha) for more information. + + torch_compile (bool): Applies only to the `"flex-fna"` backend. Whether or not to JIT + compile the attention kernel. Due to this being an experimental feature in PyTorch, we + do not recommend it, and it is guarded by context flags. Read more in + [Flex Attention + `torch.compile`](context.md#flex-attention-torchcompile). + + attention_kwargs: arguments to the [attention][natten.attention] operator, if used to + implement neighborhood cross-attention, or self attention as a fast path for + neighborhood attention. + + If `additional_{keys,values}` are specified, NATTEN usually performs a separate + cross-attention using our [attention][natten.attention] operator, and + [merges][natten.merge_attentions] the results. + + If for a given use case, the neighborhood attention problem is equivalent to self + attention (not causal along any dims, `kernel_size == (X, Y, Z)`), NATTEN will also + attempt to directly use [attention][natten.attention]. + + You can override arguments to [attention][natten.attention] by passing a + dictionary here. + + !!! example + ```python + out = na3d( + q, k, v, kernel_size=kernel_size, + ..., + attention_kwargs={ + "backend": "blackwell-fmha", + "run_persistent_kernel": True, + } + ) + ``` + + return_lse (bool): Whether or not to return the `logsumexp` tensor. `logsumexp` can be used + in the backward pass, and for [attention merging][natten.merge_attentions]. + + Returns: + output (Tensor): 6-D output tensor, with the heads last layout + (`[batch, X, Y, Z, heads, head_dim_v]`). + + logsumexp (Tensor): only returned when `return_lse=True`. 5-D logsumexp tensor, with the + heads last layout (`[batch, X, Y, Z, heads]`). + """ + return neighborhood_attention_generic( + query=query, + key=key, + value=value, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + scale=scale, + additional_keys=additional_keys, + additional_values=additional_values, + attention_kwargs=attention_kwargs, + backend=backend, + q_tile_shape=q_tile_shape, + kv_tile_shape=kv_tile_shape, + backward_q_tile_shape=backward_q_tile_shape, + backward_kv_tile_shape=backward_kv_tile_shape, + backward_kv_splits=backward_kv_splits, + backward_use_pt_reduction=backward_use_pt_reduction, + run_persistent_kernel=run_persistent_kernel, + kernel_schedule=kernel_schedule, + torch_compile=torch_compile, + return_lse=return_lse, + ) diff --git a/natten/torch-ext/natten/modules.py b/natten/torch-ext/natten/modules.py new file mode 100644 index 00000000..c3f6a59e --- /dev/null +++ b/natten/torch-ext/natten/modules.py @@ -0,0 +1,449 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# +from typing import Optional + +import torch # noqa: F401 +from torch import nn, Tensor + +from .functional import neighborhood_attention_generic +from .types import ( + CausalArg1DTypeOrDed, + CausalArg2DTypeOrDed, + CausalArg3DTypeOrDed, + CausalArgTypeOrDed, + Dimension1DTypeOrDed, + Dimension2DTypeOrDed, + Dimension3DTypeOrDed, + DimensionTypeOrDed, +) +from .utils.checks import check_all_args + + +class NeighborhoodAttentionGeneric(nn.Module): + def __init__( + self, + na_dim: int, + embed_dim: int, + num_heads: int, + kernel_size: DimensionTypeOrDed, + stride: DimensionTypeOrDed = 1, + dilation: DimensionTypeOrDed = 1, + is_causal: CausalArgTypeOrDed = False, + qkv_bias: bool = True, + qk_scale: Optional[float] = None, + proj_drop: float = 0.0, + ): + super().__init__() + kernel_size, stride, dilation, is_causal = check_all_args( + na_dim, kernel_size, stride, dilation, is_causal + ) + + if embed_dim % num_heads != 0: + raise ValueError( + "Number of attention heads must evenly divide embedding dimension, " + f"got {embed_dim=}, {num_heads=}." + ) + + self.na_dim = na_dim + self.embed_dim = embed_dim + self.num_heads = num_heads + self.head_dim = self.embed_dim // self.num_heads + self.scale = qk_scale or self.head_dim**-0.5 + self.kernel_size = kernel_size + self.stride = stride + self.dilation = dilation + self.is_causal = is_causal + + self.expected_input_tensor_rank = self.na_dim + 2 # batch, embedding dim + + self.qkv = nn.Linear(self.embed_dim, self.embed_dim * 3, bias=qkv_bias) + self.proj = nn.Linear(self.embed_dim, self.embed_dim) + self.proj_drop = nn.Dropout(proj_drop) + + def forward(self, x: Tensor) -> Tensor: + if x.dim() != self.expected_input_tensor_rank: + raise ValueError( + f"NeighborhoodAttention{self.na_dim}D expected a tensor with rank " + f"{self.expected_input_tensor_rank} ({self.na_dim} for token layout, 1 for batch, " + f"1 for embedding dimension), got {x.dim()=}." + ) + + B, *input_shape, C = x.shape + + if C != self.embed_dim: + raise ValueError( + f"Expected embedding dimension {self.embed_dim}, got {C} ({x.shape=})." + ) + + # 3, batch, *input_shape, heads, head_dim + permutation = ( + [self.na_dim + 1, 0] + + [x + 1 for x in range(self.na_dim)] + + [self.na_dim + 2, self.na_dim + 3] + ) + qkv = ( + self.qkv(x) + .reshape(B, *input_shape, 3, self.num_heads, self.head_dim) + .permute(*permutation) + ) + q, k, v = qkv[0], qkv[1], qkv[2] + x = neighborhood_attention_generic( # type: ignore[assignment] + q, + k, + v, + kernel_size=self.kernel_size, + stride=self.stride, + dilation=self.dilation, + is_causal=self.is_causal, + scale=self.scale, + ) + x = x.reshape(B, *input_shape, C) + + return self.proj_drop(self.proj(x)) + + def extra_repr(self) -> str: + return ( + f"head_dim={self.head_dim}, num_heads={self.num_heads}, " + + f"kernel_size={self.kernel_size}, " + + f"stride={self.stride}, " + + f"dilation={self.dilation}, " + + f"is_causal={self.is_causal}" + ) + + +class NeighborhoodAttention1D(NeighborhoodAttentionGeneric): + """ + 1-D Neighborhood Attention torch module. + + Performs QKV and output linear projections in addition to the [na1d][natten.na1d] operation. + + Args: + embed_dim: Embedding dimension size (a.k.a. number of channels, latent size). + !!! note + This is not `head_dim`. It's `head_dim * num_heads`. + + num_heads: Number of attention heads. + + kernel_size (Tuple[int] | int): Neighborhood window (kernel) size. + + !!! note + `kernel_size` must be smaller than or equal to `seqlen`. + + stride (Tuple[int] | int): Sliding window step size. Defaults to `1` (standard sliding + window). + + !!! note + `stride` must be smaller than or equal to `kernel_size`. + When `stride == kernel_size`, there will be no overlap between sliding windows, + which is equivalent to blocked attention (a.k.a. + [window self attention](https://arxiv.org/abs/2103.14030)). + + dilation (Tuple[int] | int): Dilation step size. Defaults to `1` (standard sliding window). + + !!! note + The product of `dilation` and `kernel_size` must be smaller than or equal to + `seqlen`. + + is_causal (Tuple[bool] | bool): Toggle causal masking. Defaults to `False` + (bi-directional). + + qkv_bias: Enable bias in the QKV linear projection. + + qk_scale: Attention scale. Defaults to `head_dim ** -0.5`. + + proj_drop: Dropout score for projection layer. Defaults is `0.0` (no dropout). + + Example: + ```python3 + import torch + from . import NeighborhoodAttention1D + + num_heads = 4 + head_dim = 128 + embed_dim = num_heads * head_dim + + model = NeighborhoodAttention1D( + embed_dim=embed_dim, + num_heads=num_heads, + kernel_size=2048, + stride=2, + dilation=4, + is_causal=True + ) + + batch = 1 + seqlen = 4096 # (1)! + + x = torch.randn(batch, seqlen, embed_dim) # (2)! + y = model(x) # (3)! + ``` + + 1. Tokens are arranged in a sequential layout of size 4096, to which we apply a + kernel size of 2048, stride 2, dilation 4, and apply causal masking. + + 2. `x.shape == [1, 4096, 512]` + 3. `y.shape == [1, 4096, 512]` + """ + + def __init__( + self, + embed_dim: int, + num_heads: int, + kernel_size: Dimension1DTypeOrDed, + stride: Dimension1DTypeOrDed = 1, + dilation: Dimension1DTypeOrDed = 1, + is_causal: CausalArg1DTypeOrDed = False, + qkv_bias: bool = True, + qk_scale: Optional[float] = None, + proj_drop: float = 0.0, + ): + super().__init__( + na_dim=1, + embed_dim=embed_dim, + num_heads=num_heads, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + proj_drop=proj_drop, + ) + + +class NeighborhoodAttention2D(NeighborhoodAttentionGeneric): + """ + 2-D Neighborhood Attention torch module. + + Performs QKV and output linear projections in addition to the [na2d][natten.na2d] operation. + + Args: + embed_dim: Embedding dimension size (a.k.a. number of channels, latent size). + !!! note + This is not `head_dim`. It's `head_dim * num_heads`. + + num_heads: Number of attention heads. + + kernel_size (Tuple[int, int] | int): Neighborhood window (kernel) size/shape. If an + integer, it will be repeated for all 2 dimensions. For example `kernel_size=3` is + reinterpreted as `kernel_size=(3, 3)`. + + !!! note + `kernel_size` must be smaller than or equal to token layout shape (`(X, Y)`) along + every dimension. + + stride (Tuple[int, int] | int): Sliding window step size/shape. Defaults to `1` (standard + sliding window). If an integer, it will be repeated for all 2 dimensions. For example + `stride=2` is reinterpreted as `stride=(2, 2)`. + + !!! note + `stride` must be smaller than or equal to `kernel_size` along every dimension. + When `stride == kernel_size`, there will be no overlap between sliding windows, + which is equivalent to blocked attention (a.k.a. + [window self attention](https://arxiv.org/abs/2103.14030)). + + dilation (Tuple[int, int] | int): Dilation step size/shape. Defaults to `1` (standard + sliding window). If an integer, it will be repeated for all 2 dimensions. For example + `dilation=4` is reinterpreted as `dilation=(4, 4)`. + + !!! note + The product of `dilation` and `kernel_size` must be smaller than or equal to + token layout shape (`(X, Y)`) along every dimension. + + is_causal (Tuple[bool, bool] | bool): Toggle causal masking. Defaults to `False` + (bi-directional). If a boolean, it will be repeated for all 2 dimensions. For example + `is_causal=True` is reinterpreted as `is_causal=(True, True)`. + + qkv_bias: Enable bias in the QKV linear projection. + + qk_scale: Attention scale. Defaults to `head_dim ** -0.5`. + + proj_drop: Dropout score for projection layer. Defaults is `0.0` (no dropout). + + Example: + ```python3 + import torch + from . import NeighborhoodAttention2D + + num_heads = 4 + head_dim = 128 + embed_dim = num_heads * head_dim + + model = NeighborhoodAttention2D( + embed_dim=embed_dim, + num_heads=num_heads, + kernel_size=(8, 16), + stride=(1, 2), + dilation=(2, 1), + is_causal=False + ) + + batch = 1 + token_layout_shape = (16, 32) # (1)! + + x = torch.randn(batch, *token_layout_shape, embed_dim) # (2)! + y = model(x) # (3)! + ``` + + 1. Tokens are arranged in a 16 x 32 layout, to which we apply a + kernel size of 8 x 16, + stride 1 x 2, + and dilation 2 x 1. + + 2. `x.shape == [1, 16, 32, 512]` + 3. `y.shape == [1, 16, 32, 512]` + """ + + def __init__( + self, + embed_dim: int, + num_heads: int, + kernel_size: Dimension2DTypeOrDed, + stride: Dimension2DTypeOrDed = 1, + dilation: Dimension2DTypeOrDed = 1, + is_causal: CausalArg2DTypeOrDed = False, + qkv_bias: bool = True, + qk_scale: Optional[float] = None, + proj_drop: float = 0.0, + ): + super().__init__( + na_dim=2, + embed_dim=embed_dim, + num_heads=num_heads, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + proj_drop=proj_drop, + ) + + +class NeighborhoodAttention3D(NeighborhoodAttentionGeneric): + """ + 3-D Neighborhood Attention torch module. + + Performs QKV and output linear projections in addition to the [na3d][natten.na3d] operation. + + Args: + embed_dim: Embedding dimension size (a.k.a. number of channels, latent size). + !!! note + This is not `head_dim`. It's `head_dim * num_heads`. + + num_heads: Number of attention heads. + + kernel_size (Tuple[int, int, int] | int): Neighborhood window (kernel) size/shape. If an + integer, it will be repeated for all 3 dimensions. For example `kernel_size=3` is + reinterpreted as `kernel_size=(3, 3, 3)`. + + !!! note + `kernel_size` must be smaller than or equal to token layout shape (`(X, Y, Z)`) + along every dimension. + + stride (Tuple[int, int, int] | int): Sliding window step size/shape. Defaults to `1` + (standard sliding window). If an integer, it will be repeated for all 3 dimensions. + For example `stride=2` is reinterpreted as `stride=(2, 2, 2)`. + + !!! note + `stride` must be smaller than or equal to `kernel_size` along every dimension. + When `stride == kernel_size`, there will be no overlap between sliding windows, + which is equivalent to blocked attention (a.k.a. + [window self attention](https://arxiv.org/abs/2103.14030)). + + dilation (Tuple[int, int, int] | int): Dilation step size/shape. Defaults to `1` (standard + sliding window). If an integer, it will be repeated for all 3 dimensions. For example + `dilation=4` is reinterpreted as `dilation=(4, 4, 4)`. + + !!! note + The product of `dilation` and `kernel_size` must be smaller than or equal to + token layout shape (`(X, Y, Z)`) along every dimension. + + is_causal (Tuple[bool, bool, bool] | bool): Toggle causal masking. Defaults to `False` + (bi-directional). If a boolean, it will be repeated for all 3 dimensions. For example + `is_causal=True` is reinterpreted as `is_causal=(True, True, True)`. + + qkv_bias: Enable bias in the QKV linear projection. + + qk_scale: Attention scale. Defaults to `head_dim ** -0.5`. + + proj_drop: Dropout score for projection layer. Defaults is `0.0` (no dropout). + + Example: + ```python3 + import torch + from . import NeighborhoodAttention3D + + num_heads = 4 + head_dim = 128 + embed_dim = num_heads * head_dim + + model = NeighborhoodAttention3D( + embed_dim=embed_dim, + num_heads=num_heads, + kernel_size=(4, 8, 12), + stride=(1, 1, 4), + dilation=(1, 2, 1), + is_causal=(True, False, False) + ) + + batch = 1 + token_layout_shape = (12, 16, 20) # (1)! + + x = torch.randn(batch, *token_layout_shape, embed_dim) # (2)! + y = model(x) # (3)! + ``` + + 1. Tokens are arranged in a 12 x 16 x 20 layout, to which we apply a + kernel size of 4 x 8 x 12, + stride 1 x 1 x 4, + dilation 1 x 2 x 1, and apply causal masking to the left-most dimension (12). + + 2. `x.shape == [1, 12, 16, 20, 512]` + 3. `y.shape == [1, 12, 16, 20, 512]` + """ + + def __init__( + self, + embed_dim: int, + num_heads: int, + kernel_size: Dimension3DTypeOrDed, + stride: Dimension3DTypeOrDed = 1, + dilation: Dimension3DTypeOrDed = 1, + is_causal: CausalArg3DTypeOrDed = False, + qkv_bias: bool = True, + qk_scale: Optional[float] = None, + proj_drop: float = 0.0, + ): + super().__init__( + na_dim=3, + embed_dim=embed_dim, + num_heads=num_heads, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + is_causal=is_causal, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + proj_drop=proj_drop, + ) diff --git a/natten/torch-ext/natten/token_permute/__init__.py b/natten/torch-ext/natten/token_permute/__init__.py new file mode 100644 index 00000000..381de37e --- /dev/null +++ b/natten/torch-ext/natten/token_permute/__init__.py @@ -0,0 +1,32 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +from ..token_permute.frontend import ( + token_permute_operation, + token_unpermute_operation, +) + +__all__ = [ + "token_permute_operation", + "token_unpermute_operation", +] diff --git a/natten/torch-ext/natten/token_permute/cutlass_impl.py b/natten/torch-ext/natten/token_permute/cutlass_impl.py new file mode 100644 index 00000000..21694256 --- /dev/null +++ b/natten/torch-ext/natten/token_permute/cutlass_impl.py @@ -0,0 +1,286 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +import functools +from typing import Tuple + +import torch +from torch import Tensor +from torch.amp import custom_bwd, custom_fwd +from torch.autograd import Function + +amp_fwd = functools.partial(custom_fwd, device_type="cuda") +amp_bwd = functools.partial(custom_bwd, device_type="cuda") + +from .._libnatten import ( + HAS_LIBNATTEN, + token_permute_1d, + token_permute_2d, + token_permute_3d, + token_unpermute_1d, + token_unpermute_2d, + token_unpermute_3d, +) +from ..types import DimensionType, NoneType +from ..utils import log +from ..utils.device import get_device_cc, is_cuda + +logger = log.get_logger(__name__) + + +def can_run_cutlass_tokperm(tensor: Tensor) -> bool: + if not HAS_LIBNATTEN: + logger.debug( + "Can't use libnatten TokPerm kernels, because libnatten is not available." + ) + return False + + if not is_cuda(tensor.device): + logger.debug( + "Can't use libnatten TokPerm kernels, because input is not a CUDA tensor." + ) + return False + + is_fp8_allowed = get_device_cc(tensor.device) in [100, 103] + if tensor.dtype not in [ + torch.float32, + torch.float16, + torch.bfloat16, + torch.float16, + ] and ( + is_fp8_allowed and tensor.dtype not in [torch.float8_e5m2, torch.float8_e4m3fn] + ): + logger.debug( + f"Can't use libnatten TokPerm kernels; unexpected dtype {tensor.dtype}." + ) + return False + + return True + + +PERMUTE_OPS = {1: token_permute_1d, 2: token_permute_2d, 3: token_permute_3d} +UNPERMUTE_OPS = {1: token_unpermute_1d, 2: token_unpermute_2d, 3: token_unpermute_3d} + + +def make_cutlass_token_permute_autograd_fn(na_dim): + assert na_dim in [1, 2, 3] + + class CutlassTokenPermuteAutogradFn(Function): + @staticmethod + @amp_fwd + def forward( + ctx, + tensor: Tensor, + tile_shape: DimensionType, + dilation: DimensionType, + flip_tiled_dims: bool, + ) -> Tensor: + + output = PERMUTE_OPS[na_dim]( + tensor, + tile_shape=tile_shape, + dilation=dilation, + flip_tiled_dims=flip_tiled_dims, + ) + + ctx.tile_shape = tile_shape + ctx.dilation = dilation + ctx.flip_tiled_dims = flip_tiled_dims + ctx.token_layout = tuple(x for x in tensor.shape[1 : na_dim + 1]) + assert len(ctx.token_layout) == na_dim + + return output + + @staticmethod + @amp_bwd + def backward(ctx, d_output: Tensor) -> Tuple[ + Tensor, + NoneType, + NoneType, + NoneType, + ]: + + d_output_unpermuted = UNPERMUTE_OPS[na_dim]( + d_output, + token_layout_shape=ctx.token_layout, + tile_shape=ctx.tile_shape, + dilation=ctx.dilation, + flip_tiled_dims=ctx.flip_tiled_dims, + ) + + return ( + d_output_unpermuted, + None, + None, + None, + ) + + return CutlassTokenPermuteAutogradFn + + +def make_cutlass_token_unpermute_autograd_fn(na_dim): + assert na_dim in [1, 2, 3] + + class CutlassTokenUnPermuteAutogradFn(Function): + @staticmethod + @amp_fwd + def forward( + ctx, + tensor: Tensor, + token_layout: DimensionType, + tile_shape: DimensionType, + dilation: DimensionType, + flip_tiled_dims: bool, + ) -> Tensor: + + output = UNPERMUTE_OPS[na_dim]( + tensor, + token_layout_shape=token_layout, + tile_shape=tile_shape, + dilation=dilation, + flip_tiled_dims=flip_tiled_dims, + ) + + ctx.tile_shape = tile_shape + ctx.dilation = dilation + ctx.flip_tiled_dims = flip_tiled_dims + + return output + + @staticmethod + @amp_bwd + def backward(ctx, d_output: Tensor) -> Tuple[ + Tensor, + NoneType, + NoneType, + NoneType, + NoneType, + ]: + + d_output_permuted = PERMUTE_OPS[na_dim]( + d_output, + tile_shape=ctx.tile_shape, + dilation=ctx.dilation, + flip_tiled_dims=ctx.flip_tiled_dims, + ) + + return ( + d_output_permuted, + None, + None, + None, + None, + ) + + return CutlassTokenUnPermuteAutogradFn + + +CutlassTokenPermute1DAutogradFn = make_cutlass_token_permute_autograd_fn(1) +CutlassTokenPermute2DAutogradFn = make_cutlass_token_permute_autograd_fn(2) +CutlassTokenPermute3DAutogradFn = make_cutlass_token_permute_autograd_fn(3) + +CutlassTokenUnPermute1DAutogradFn = make_cutlass_token_unpermute_autograd_fn(1) +CutlassTokenUnPermute2DAutogradFn = make_cutlass_token_unpermute_autograd_fn(2) +CutlassTokenUnPermute3DAutogradFn = make_cutlass_token_unpermute_autograd_fn(3) + +CutlassTokenPermuteAutogradFns = { + 1: CutlassTokenPermute1DAutogradFn, + 2: CutlassTokenPermute2DAutogradFn, + 3: CutlassTokenPermute3DAutogradFn, +} + +CutlassTokenUnPermuteAutogradFns = { + 1: CutlassTokenUnPermute1DAutogradFn, + 2: CutlassTokenUnPermute2DAutogradFn, + 3: CutlassTokenUnPermute3DAutogradFn, +} + + +def token_permute_cutlass( + tensor: Tensor, + tile_shape: DimensionType, + dilation: DimensionType, + flip_tiled_dims, +) -> Tensor: + if tensor.dim() not in [4, 5, 6]: + raise ValueError( + "Expected 4D, 5D, or 6D tensor (corresponding to NA1D, 2D, 3D), " + f"got {tensor.dim()}D input." + ) + + na_dim = tensor.dim() - 3 + assert na_dim in [1, 2, 3] + + if len(tile_shape) != na_dim: + raise ValueError( + f"Expected {na_dim}D tiler for NA{na_dim}D, " f"got {tile_shape=}." + ) + + if not can_run_cutlass_tokperm(tensor): + raise NotImplementedError( + "Use case is not compatible with CUTLASS Token Permute." + ) + + tensor = tensor.contiguous() + output = CutlassTokenPermuteAutogradFns[na_dim].apply( + tensor, + tile_shape, + dilation, + flip_tiled_dims, + ) + + return output + + +def token_unpermute_cutlass( + tensor: Tensor, + token_layout_shape: DimensionType, + tile_shape: DimensionType, + dilation: DimensionType, + flip_tiled_dims: bool, +) -> Tensor: + if tensor.dim() != 4: + raise ValueError(f"Expected flattened 4D tensor, got {tensor.dim()}D input.") + + na_dim = len(token_layout_shape) + + if len(tile_shape) != na_dim: + raise ValueError( + f"Expected {na_dim}D tiler for NA{na_dim}D, " f"got {tile_shape=}." + ) + + if not can_run_cutlass_tokperm(tensor): + raise NotImplementedError( + "Use case is not compatible with CUTLASS Token UnPermute." + ) + + tensor = tensor.contiguous() + output = CutlassTokenUnPermuteAutogradFns[na_dim].apply( + tensor, + token_layout_shape, + tile_shape, + dilation, + flip_tiled_dims, + ) + + return output diff --git a/natten/torch-ext/natten/token_permute/frontend.py b/natten/torch-ext/natten/token_permute/frontend.py new file mode 100644 index 00000000..52d3183f --- /dev/null +++ b/natten/torch-ext/natten/token_permute/frontend.py @@ -0,0 +1,137 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +from typing import Optional + +from torch import Tensor + +from .._environment import USE_TORCH_IMPL_DEFAULT +from ..token_permute.cutlass_impl import ( + can_run_cutlass_tokperm, + token_permute_cutlass, + token_unpermute_cutlass, +) +from ..token_permute.torch_impl import token_permute_torch, token_unpermute_torch +from ..types import DimensionType +from ..utils import log +from ..utils.tuples import ceil_div_tuple, mul_tuple + +logger = log.get_logger(__name__) + + +def token_permute_operation( + tensor: Tensor, + tile_shape: DimensionType, + dilation: Optional[DimensionType] = None, + flip_tiled_dims: bool = True, + use_torch: bool = USE_TORCH_IMPL_DEFAULT, +) -> tuple[Tensor, DimensionType, DimensionType]: + if tensor.dim() not in [4, 5, 6]: + raise ValueError( + "Expected 4D, 5D, or 6D tensor (corresponding to NA1D, 2D, 3D), " + f"got {tensor.dim()}D input." + ) + + na_dim = tensor.dim() - 3 + assert na_dim in [1, 2, 3] + + if len(tile_shape) != na_dim: + raise ValueError( + f"Expected {na_dim}D tiler for NA{na_dim}D, " f"got {tile_shape=}." + ) + + if dilation is not None and len(dilation) != na_dim: + raise ValueError( + f"Expected {na_dim}D dilation for NA{na_dim}D, " f"got {dilation=}." + ) + + dilation_: DimensionType = dilation or tuple(1 for _ in range(na_dim)) # type: ignore[assignment] + + tensor = tensor.contiguous() + batch, *token_layout_, heads, dim = tensor.shape + token_layout: DimensionType = tuple(x for x in token_layout_) # type: ignore[assignment] + + token_layout_post_dilation: DimensionType = mul_tuple(ceil_div_tuple(ceil_div_tuple(token_layout, tile_shape), dilation_), tile_shape) # type: ignore[assignment] + + if not use_torch and can_run_cutlass_tokperm(tensor): + output = token_permute_cutlass( + tensor, + tile_shape=tile_shape, + dilation=dilation_, + flip_tiled_dims=flip_tiled_dims, + ) + else: + output = token_permute_torch( + tensor, + tile_shape=tile_shape, + dilation=dilation_, + flip_tiled_dims=flip_tiled_dims, + ) + + return output, token_layout, token_layout_post_dilation + + +def token_unpermute_operation( + tensor: Tensor, + token_layout_shape: DimensionType, + tile_shape: DimensionType, + dilation: Optional[DimensionType] = None, + flip_tiled_dims: bool = True, + use_torch: bool = USE_TORCH_IMPL_DEFAULT, +) -> Tensor: + if tensor.dim() != 4: + raise ValueError(f"Expected flattened 4D tensor, got {tensor.dim()}D input.") + + na_dim = len(token_layout_shape) + + if len(tile_shape) != na_dim: + raise ValueError( + f"Expected {na_dim}D tiler for NA{na_dim}D, " f"got {tile_shape=}." + ) + + if dilation is not None and len(dilation) != na_dim: + raise ValueError( + f"Expected {na_dim}D dilation for NA{na_dim}D, " f"got {dilation=}." + ) + + dilation_: DimensionType = dilation or tuple(1 for _ in range(na_dim)) # type: ignore[assignment] + + tensor = tensor.contiguous() + if not use_torch and can_run_cutlass_tokperm(tensor): + output = token_unpermute_cutlass( + tensor, + token_layout_shape, + tile_shape=tile_shape, + dilation=dilation_, + flip_tiled_dims=flip_tiled_dims, + ) + else: + output = token_unpermute_torch( + tensor, + token_layout_shape, + tile_shape=tile_shape, + dilation=dilation_, + flip_tiled_dims=flip_tiled_dims, + ) + + return output diff --git a/natten/torch-ext/natten/token_permute/torch_impl.py b/natten/torch-ext/natten/token_permute/torch_impl.py new file mode 100644 index 00000000..b0ba3dce --- /dev/null +++ b/natten/torch-ext/natten/token_permute/torch_impl.py @@ -0,0 +1,368 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# +import math + +import torch +from torch import Tensor + +from ..types import DimensionType +from ..utils import log +from ..utils.environment import is_torch_compiling +from ..utils.tuples import ceil_div_tuple, mul_tuple, sub_tuple + +logger = log.get_logger(__name__) + + +DISABLE_PADDING_WARNING = True +TOKEN_PERMUTE_PADDING_RATIO_LIMIT_UNTIL_WARNING = 0.5 + + +def _maybe_pad( + tensor: Tensor, tile_shape: DimensionType, dilation: DimensionType +) -> Tensor: + if tensor.dim() not in [4, 5, 6]: + raise ValueError( + "Expected 4D, 5D, or 6D tensor (corresponding to NA1D, 2D, 3D), " + f"got {tensor.dim()}D input." + ) + + na_dim = tensor.dim() - 3 + assert na_dim in [1, 2, 3] + + if len(tile_shape) != na_dim: + raise ValueError( + f"Expected {na_dim}D tiler for NA{na_dim}D, " f"got {tile_shape=}." + ) + + if dilation is not None and len(dilation) != na_dim: + raise ValueError( + f"Expected {na_dim}D dilation for NA{na_dim}D, " f"got {dilation=}." + ) + + token_layout = tensor.shape[1 : na_dim + 1] + tile_shape_ = tuple(x for x in tile_shape) + if dilation is not None: + # NOTE: LCM? + # tile_shape_ = tuple(math.lcm(t, d) for t, d in zip(tile_shape, dilation)) + tile_shape_ = tuple(t * d for t, d in zip(tile_shape, dilation)) + + rest = tuple((x + t - 1) // t for x, t in zip(token_layout, tile_shape_)) + residual = tuple(r * t - x for x, t, r in zip(token_layout, tile_shape_, rest)) + + assert all(res >= 0 for res in residual) + + if not DISABLE_PADDING_WARNING and any( + res / sz > TOKEN_PERMUTE_PADDING_RATIO_LIMIT_UNTIL_WARNING + for res, sz in zip(residual, token_layout) + ): + padded_token_layout = tuple(x + p for x, p in zip(token_layout, residual)) + logger.warning( + "Potentially excessive padding detected in token permute: " + f"input shape {token_layout} will be padded to {padded_token_layout} to handle " + "token permutation, which can result in excessive memory usage, and " + "performance implications. Consider choosing your tile shapes, input shapes " + "(and dilation if you use it) accordingly. Refer to NATTEN docs for more info." + ) + + if any(res > 0 for res in residual): + padding = [0, 0, 0, 0] # head_dim_left, head_dim_right, heads_left, heads_right + for res in reversed(residual): + padding.append(0) # left pad + padding.append(res) # right pad + tensor_padded = torch.nn.functional.pad(tensor, padding, "constant", 0) + else: + tensor_padded = tensor + + return tensor_padded + + +def _token_permute( + tensor: Tensor, + tile_shape: DimensionType, + dilation: DimensionType, + flip_tiled_dims: bool, +) -> Tensor: + if tensor.dim() not in [4, 5, 6]: + raise ValueError( + "Expected 4D, 5D, or 6D tensor (corresponding to NA1D, 2D, 3D), " + f"got {tensor.dim()}D input." + ) + + na_dim = tensor.dim() - 3 + assert na_dim in [1, 2, 3] + + if len(tile_shape) != na_dim: + raise ValueError( + f"Expected {na_dim}D tiler for NA{na_dim}D, " f"got {tile_shape=}." + ) + + dilation = dilation or tuple(1 for _ in range(na_dim)) + + if len(dilation) != na_dim: + raise ValueError( + f"Expected {na_dim}D dilation for NA{na_dim}D, " f"got {dilation=}." + ) + + batch, *token_layout, heads, dim = tensor.shape + + if any( + x % d != 0 or (x // d) % t != 0 + for x, t, d in zip(token_layout, tile_shape, dilation) + ): + raise ValueError( + "Tensor must be divisible by static tile shape and dilation, but got " + f"{tensor.shape=}, {tile_shape=}, {dilation=}." + ) + + num_dilation_groups = math.prod(dilation) + token_layout_post_dilation = tuple(x // d for x, d in zip(token_layout, dilation)) + rest = tuple(x // d // t for x, t, d in zip(token_layout, tile_shape, dilation)) + logical_divide_dims = [] + for d, r, t in zip(dilation, rest, tile_shape): + logical_divide_dims += [r, t, d] + + # Two permutations at once: + # 1. logical divide to tiled divide + # 2. (optionally) flip order of tiled modes (i.e. (X,Y,Z) -> (Z,Y,X)) for compatibility with + # CuTe's identity layout mapping. + permutation_idxes_r = [] + permutation_idxes_t = [] + permutation_idxes_d = [] + for i in range(na_dim): + if flip_tiled_dims: + permutation_idxes_r += [(na_dim - i - 1) * 3 + 1] + permutation_idxes_t += [(na_dim - i - 1) * 3 + 2] + permutation_idxes_d += [(na_dim - i - 1) * 3 + 3] + else: + permutation_idxes_r += [i * 3 + 1] + permutation_idxes_t += [i * 3 + 2] + permutation_idxes_d += [i * 3 + 3] + + permutation_idxes = ( + [0] + + permutation_idxes_d + + permutation_idxes_r + + permutation_idxes_t + + [na_dim * 3 + 1, na_dim * 3 + 2] + ) + + # View, not copy + tensor_tiled = tensor.view(batch, *logical_divide_dims, heads, dim) + if not is_torch_compiling(): + assert tensor_tiled.data_ptr() == tensor.data_ptr() + + # View, not copy + tensor_permuted = tensor_tiled.permute(*permutation_idxes) + if not is_torch_compiling(): + assert tensor_permuted.data_ptr() == tensor_tiled.data_ptr() + + # Reshape back and copy + tensor_flatten = tensor_permuted.reshape( + num_dilation_groups * batch, math.prod(token_layout_post_dilation), heads, dim + ).contiguous() + # NOTE: token permute without dilation is a no-op for 1-D + # assert na_dim == 1 or tensor_flatten.data_ptr() != tensor_permuted.data_ptr() + assert tensor_flatten.is_contiguous() + + return tensor_flatten + + +def _token_unpermute( + tensor: Tensor, + token_layout: DimensionType, + tile_shape: DimensionType, + dilation: DimensionType, + flip_tiled_dims: bool, +): + if tensor.dim() != 4: + raise ValueError(f"Expected flattened 4D tensor, got {tensor.dim()}D input.") + + na_dim = len(token_layout) + assert na_dim in [1, 2, 3] + + if len(tile_shape) != na_dim: + raise ValueError( + f"Expected {na_dim}D tiler for NA{na_dim}D, " f"got {tile_shape=}." + ) + + dilation = dilation or tuple(1 for _ in range(na_dim)) + + if len(dilation) != na_dim: + raise ValueError( + f"Expected {na_dim}D dilation for NA{na_dim}D, " f"got {dilation=}." + ) + + num_dilation_groups = math.prod(dilation) + + batch, seqlen, heads, dim = tensor.shape + + if batch % num_dilation_groups != 0: + raise ValueError( + "Expected batch size in token-permuted tensor to be divisible by " + f"number of dilation groups {num_dilation_groups} ({dilation=}), got {batch=}." + ) + + batch_actual = batch // num_dilation_groups + + rest_shape = ceil_div_tuple(ceil_div_tuple(token_layout, tile_shape), dilation) + token_layout_padded = mul_tuple(mul_tuple(rest_shape, tile_shape), dilation) + + # View, not copy + rest_shape_ = reversed(rest_shape) if flip_tiled_dims else rest_shape + tile_shape_ = reversed(tile_shape) if flip_tiled_dims else tile_shape + dilation_ = reversed(dilation) if flip_tiled_dims else dilation + tensor_tiled = tensor.view( + batch_actual, *dilation_, *rest_shape_, *tile_shape_, heads, dim + ) + if not is_torch_compiling(): + assert tensor_tiled.data_ptr() == tensor.data_ptr() + + # Undo permutation + # batch + permutation_idxes = [0] + + # dilation, rest, tile -> rest, tile, dilation + for i in range(na_dim): + if flip_tiled_dims: + permutation_idxes += [2 * na_dim - i, 3 * na_dim - i, na_dim - i] + else: + permutation_idxes += [na_dim + i + 1, 2 * na_dim + i + 1, i + 1] + + # heads, head_dim + permutation_idxes += [na_dim * 3 + 1, na_dim * 3 + 2] + + # View, not copy + tensor_permuted = tensor_tiled.permute(*permutation_idxes) + if not is_torch_compiling(): + assert tensor_permuted.data_ptr() == tensor_tiled.data_ptr() + + # Reshape back and copy + out = tensor_permuted.reshape( + batch_actual, *token_layout_padded, heads, dim + ).contiguous() + # NOTE: token permute without dilation is a no-op for 1-D + # assert na_dim == 1 or out.data_ptr() != tensor_permuted.data_ptr() + assert out.is_contiguous() + + return out + + +def _maybe_unpad(tensor: Tensor, padding: DimensionType): + if tensor.dim() not in [4, 5, 6]: + raise ValueError( + "Expected 4D, 5D, or 6D tensor (corresponding to NA1D, 2D, 3D), " + f"got {tensor.dim()}D input." + ) + + na_dim = tensor.dim() - 3 + assert na_dim in [1, 2, 3] + + if len(padding) != na_dim: + raise ValueError( + f"Expected {na_dim}D padding shape for NA{na_dim}D, " f"got {padding=}." + ) + + token_layout = tensor.shape[1 : na_dim + 1] + + # Slice + if any(p for p in padding): + assert all(p >= 0 for p in padding) + + orig_lens = tuple(x - p for x, p in zip(token_layout, padding)) + + # TODO: there must be a better way + if len(orig_lens) == 1: + x = orig_lens[0] + return tensor[:, :x].contiguous() + elif len(orig_lens) == 2: + x, y = orig_lens + return tensor[:, :x, :y].contiguous() + elif len(orig_lens) == 3: + x, y, z = orig_lens + return tensor[:, :x, :y, :z].contiguous() + else: + raise NotImplementedError() + + return tensor + + +def token_permute_torch( + tensor: Tensor, + tile_shape: DimensionType, + dilation: DimensionType, + flip_tiled_dims: bool, +) -> Tensor: + if tensor.dim() not in [4, 5, 6]: + raise ValueError( + "Expected 4D, 5D, or 6D tensor (corresponding to NA1D, 2D, 3D), " + f"got {tensor.dim()}D input." + ) + + tensor_pad = _maybe_pad(tensor, tile_shape=tile_shape, dilation=dilation) + output = _token_permute( + tensor_pad, + tile_shape=tile_shape, + dilation=dilation, + flip_tiled_dims=flip_tiled_dims, + ) + + return output + + +def token_unpermute_torch( + tensor: Tensor, + token_layout: DimensionType, + tile_shape: DimensionType, + dilation: DimensionType, + flip_tiled_dims: bool, +) -> Tensor: + if tensor.dim() != 4: + raise ValueError(f"Expected flattened 4D tensor, got {tensor.dim()}D input.") + + token_layout_padded = mul_tuple( + mul_tuple( + ceil_div_tuple(ceil_div_tuple(token_layout, tile_shape), dilation), + dilation, + ), + tile_shape, + ) + padding = sub_tuple(token_layout_padded, token_layout) + + output = _maybe_unpad( + _token_unpermute( + tensor, + token_layout=token_layout, + tile_shape=tile_shape, + dilation=dilation, + flip_tiled_dims=flip_tiled_dims, + ), + padding=padding, + ) + + return output + + +__all__ = [ + "token_permute_torch", + "token_unpermute_torch", +] diff --git a/natten/torch-ext/natten/types.py b/natten/torch-ext/natten/types.py new file mode 100644 index 00000000..80498090 --- /dev/null +++ b/natten/torch-ext/natten/types.py @@ -0,0 +1,85 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +from enum import Enum +from typing import Tuple, Union + +NoneType = type(None) + +Dimension1DType = Tuple[int] +Dimension2DType = Tuple[int, int] +Dimension3DType = Tuple[int, int, int] + +CausalArg1DType = Tuple[bool] +CausalArg2DType = Tuple[bool, bool] +CausalArg3DType = Tuple[bool, bool, bool] + +# NOTE: switch to | when < 3.10 support is dropped +Dimension1DTypeOrDed = Union[int, Dimension1DType] +Dimension2DTypeOrDed = Union[int, Dimension2DType] +Dimension3DTypeOrDed = Union[int, Dimension3DType] + +CausalArg1DTypeOrDed = Union[bool, CausalArg1DType] +CausalArg2DTypeOrDed = Union[bool, CausalArg2DType] +CausalArg3DTypeOrDed = Union[bool, CausalArg3DType] + +DimensionType = Union[Dimension1DType, Dimension2DType, Dimension3DType] +CausalArgType = Union[CausalArg1DType, CausalArg2DType, CausalArg3DType] + +DimensionTypeOrDed = Union[int, DimensionType] +CausalArgTypeOrDed = Union[bool, CausalArgType] + +# (query_tile_shape, kv_tile_shape) +QKTileShapeType = Union[ + Tuple[Dimension1DType, Dimension1DType], + Tuple[Dimension2DType, Dimension2DType], + Tuple[Dimension3DType, Dimension3DType], +] + + +# TODO: Only applies to Hopper FMHA/FNA for now -- extend to other applicable kernels +class KernelSchedule(Enum): + NonPersistent = 0 + WarpSpecializedCooperative = 1 + WarpSpecializedPingpong = 2 + + +CutlassFnaForwardConfigType = QKTileShapeType +CutlassFnaBackwardConfigType = QKTileShapeType +CutlassBlackwellFnaForwardConfigType = QKTileShapeType +CutlassBlackwellFnaBackwardConfigType = QKTileShapeType +CutlassHopperFnaForwardConfigType = Tuple[QKTileShapeType, KernelSchedule] +CutlassHopperFnaBackwardConfigType = QKTileShapeType +FlexFnaForwardConfigType = QKTileShapeType + +# FMHA configs +FmhaForwardConfigType = Tuple[int, int] + +CutlassFmhaForwardConfigType = FmhaForwardConfigType +CutlassFmhaBackwardConfigType = FmhaForwardConfigType + +FlexFmhaForwardConfigType = FmhaForwardConfigType +CutlassBlackwellFmhaForwardConfigType = FmhaForwardConfigType +CutlassBlackwellFmhaBackwardConfigType = FmhaForwardConfigType +CutlassHopperFmhaForwardConfigType = Tuple[FmhaForwardConfigType, KernelSchedule] +CutlassHopperFmhaBackwardConfigType = FmhaForwardConfigType diff --git a/natten/torch-ext/natten/utils/__init__.py b/natten/torch-ext/natten/utils/__init__.py new file mode 100644 index 00000000..99f6c1e8 --- /dev/null +++ b/natten/torch-ext/natten/utils/__init__.py @@ -0,0 +1,22 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# diff --git a/natten/torch-ext/natten/utils/checks.py b/natten/torch-ext/natten/utils/checks.py new file mode 100644 index 00000000..4c1e339d --- /dev/null +++ b/natten/torch-ext/natten/utils/checks.py @@ -0,0 +1,726 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +import functools +from collections.abc import Sequence +from typing import Any, Optional, Tuple, Union + +import torch # noqa: F401 +from torch import Tensor + +from ..types import CausalArgType, DimensionType, KernelSchedule, NoneType +from ..utils import log +from ..utils.environment import is_torch_compiling +from ..utils.tuples import create_causal_arg_from_bool, create_dim_from_int +from ..utils.varlen import generate_varlen_parameters + +logger = log.get_logger(__name__) + + +def log_or_raise_error( + msg: str, raise_error: bool = False, exception: Any = RuntimeError +): + if raise_error: + raise exception(msg) + else: + logger.debug(msg) + + +def _universal_tensor_checks( + query: Tensor, key: Tensor, value: Tensor, raise_error: bool = True +) -> bool: + target_fn = functools.partial(log_or_raise_error, raise_error=raise_error) + + if query.is_sparse or key.is_sparse or value.is_sparse: + target_fn( + "NATTEN does not support sparse tensors.", exception=NotImplementedError + ) + return False + + if query.is_nested or key.is_nested or value.is_nested: + target_fn( + "NATTEN does not support nested tensors.", exception=NotImplementedError + ) + return False + + if query.device != key.device or query.device != value.device: + target_fn( + "Query, key, and value must be on the same device, " + f"got {query.device=}, {key.device=}, {value.device=}.", + exception=ValueError, + ) + return False + + if query.dtype != key.dtype or query.dtype != value.dtype: + target_fn( + "Query, key, and value must assume the same data type, " + f"got {query.dtype=}, {key.dtype=}, {value.dtype=}.", + exception=ValueError, + ) + return False + + return True + + +def na_tensor_checks( + query: Tensor, + key: Tensor, + value: Tensor, + must_match_head_dims: bool = False, + supports_gqa_mqa: bool = True, + raise_error: bool = True, + backend_name: Optional[str] = None, +) -> bool: + backend_name = backend_name or "This operation/backend" + if not _universal_tensor_checks(query, key, value): + return False + + target_fn = functools.partial(log_or_raise_error, raise_error=raise_error) + + if query.dim() != key.dim() or query.dim() != value.dim(): + target_fn( + "Query, key, and value must have the same rank, " + f"got {query.dim()=}, {key.dim()=}, {value.dim()=}.", + exception=ValueError, + ) + return False + + if query.dim() not in [4, 5, 6]: + target_fn( + "Expected 4-D, 5-D, or 6-D tensors as inputs (corresponding to NA1D, NA2D, and NA3D), " + f"got {query.dim()=}.", + exception=ValueError, + ) + return False + + na_dim = query.dim() - 3 # minus batch, heads, head_dim + + if query.shape[-1] != key.shape[-1]: + target_fn( + f"Q and K head dims must match, got {query.shape[-1]=}, {key.shape[-1]=}.", + exception=ValueError, + ) + return False + + if query.shape[0] != key.shape[0] or query.shape[0] != value.shape[0]: + target_fn( + "Q, K, and V must match in batch size, got " + f"{query.shape[0]=}, {key.shape[0]=}, {value.shape[0]=}.", + exception=ValueError, + ) + return False + + if must_match_head_dims and query.shape[-1] != value.shape[-1]: + target_fn( + f"{backend_name} does not support different head dims for QK and V, got " + f"{query.shape[-1]=}, {value.shape[-1]=}.", + exception=ValueError, + ) + return False + + if ( + query.shape[1 : na_dim + 1] != key.shape[1 : na_dim + 1] + or query.shape[1 : na_dim + 1] != value.shape[1 : na_dim + 1] + ): + target_fn( + "Neighborhood Attention operations require Q, K, and V to match in their token layouts, got " + f"{query.shape[1:na_dim+1]=}, {key.shape[1:na_dim+1]=}, {value.shape[1:na_dim+1]=}.", + exception=ValueError, + ) + return False + + if not supports_gqa_mqa and ( + query.shape[-2] != key.shape[-2] or query.shape[-2] != value.shape[-2] + ): + target_fn( + f"{backend_name} does not support GQA/MQA, therefore number of heads in Q, K, and V " + f"must match, got {query.shape[-2]=}, {key.shape[-2]=}, {value.shape[-2]=}.", + exception=ValueError, + ) + return False + + if supports_gqa_mqa: + if key.shape[-2] != value.shape[-2]: + target_fn( + "Key and value must always have the same number of heads, got " + f"{key.shape[-2]=}, {value.shape[-2]=}.", + exception=ValueError, + ) + return False + + heads_q = query.shape[-2] + heads_kv = key.shape[-2] + + if heads_q < heads_kv or heads_q % heads_kv != 0: + target_fn( + "Key/value heads must evenly divide query heads, got " + f"{heads_q=}, {heads_kv=}.", + exception=ValueError, + ) + return False + + return True + + +def fmha_tensor_checks( + query: Tensor, + key: Tensor, + value: Tensor, + must_match_head_dims: bool = False, + supports_gqa_mqa: bool = True, + raise_error: bool = True, + backend_name: Optional[str] = None, +) -> bool: + backend_name = backend_name or "This operation/backend" + if not _universal_tensor_checks(query, key, value): + return False + + target_fn = functools.partial(log_or_raise_error, raise_error=raise_error) + + if query.dim() != key.dim() or query.dim() != value.dim(): + target_fn( + "Query, key, and value must have the same rank, " + f"got {query.dim()=}, {key.dim()=}, {value.dim()=}.", + exception=ValueError, + ) + return False + + if query.dim() != 4: + target_fn( + "Expected 4-D tensors as inputs to FMHA, " f"got {query.dim()=}.", + exception=ValueError, + ) + return False + + if query.shape[-1] != key.shape[-1]: + target_fn( + f"Q and K head dims must match, got {query.shape[-1]=}, {key.shape[-1]=}.", + exception=ValueError, + ) + return False + + if must_match_head_dims and query.shape[-1] != value.shape[-1]: + target_fn( + f"{backend_name} does not support different head dims for QK and V, got " + f"{query.shape[-1]=}, {value.shape[-1]=}.", + exception=ValueError, + ) + return False + + if query.shape[0] != key.shape[0] or query.shape[0] != value.shape[0]: + target_fn( + "Q, K, and V must match in batch size, got " + f"{query.shape[0]=}, {key.shape[0]=}, {value.shape[0]=}.", + exception=ValueError, + ) + return False + + if key.shape[1] != value.shape[1]: + target_fn( + f"K and V must match in sequence length, got {key.shape[1]=}, {value.shape[1]=}.", + exception=ValueError, + ) + return False + + if not supports_gqa_mqa and ( + query.shape[-2] != key.shape[-2] or query.shape[-2] != value.shape[-2] + ): + target_fn( + f"{backend_name} does not support GQA/MQA, therefore number of heads in Q, K, and V " + f"must match, got {query.shape[-2]=}, {key.shape[-2]=}, {value.shape[-2]=}.", + exception=ValueError, + ) + return False + + if supports_gqa_mqa: + if key.shape[-2] != value.shape[-2]: + target_fn( + "Key and value must always have the same number of heads, got " + f"{key.shape[-2]=}, {value.shape[-2]=}.", + exception=ValueError, + ) + return False + + heads_q = query.shape[-2] + heads_kv = key.shape[-2] + + if heads_q < heads_kv or heads_q % heads_kv != 0: + target_fn( + "Key/value heads must evenly divide query heads, got " + f"{heads_q=}, {heads_kv=}.", + exception=ValueError, + ) + return False + + return True + + +def additional_kv_tensor_checks( + query: Tensor, + key: Tensor, + value: Tensor, + add_key: Optional[Tensor] = None, + add_value: Optional[Tensor] = None, + must_match_head_dims: bool = False, + supports_gqa_mqa: bool = True, +): + + if (add_key is not None) ^ (add_value is not None): + raise ValueError( + "`additional_keys` and `additional_values` must be either both Tensors or None." + ) + + if add_key is None: + return + + assert add_key is not None and add_value is not None + + _universal_tensor_checks(query, add_key, add_value) + + if query.shape[-1] != add_key.shape[-1]: + raise ValueError( + f"Q and K head dims must match, got {query.shape[-1]=}, {add_key.shape[-1]=}." + ) + + if must_match_head_dims and query.shape[-1] != add_value.shape[-1]: + raise ValueError( + "This operation does not support different head dims for QK and V, got " + f"{query.shape[-1]=}, {add_value.shape[-1]=}." + ) + + if query.shape[0] != add_key.shape[0] or query.shape[0] != add_value.shape[0]: + raise ValueError( + "Q, additional K, and additional V must match in batch size, got " + f"{query.shape[0]=}, {add_key.shape[0]=}, {add_value.shape[0]=}." + ) + + if add_key.shape[1] != add_value.shape[1]: + raise ValueError( + f"Additional K and V must match in sequence length, got {add_key.shape[1]=}, " + f"{add_value.shape[1]=}." + ) + + if key.shape[0] != add_key.shape[0] or value.shape[0] != add_value.shape[0]: + raise ValueError( + "Additional key/value tokens must match the self attention key/value tokens in batch " + f"size, got {key.shape[0]=} != {add_key.shape[0]=}, and " + f"{value.shape[0]=} != {add_value.shape[0]=}." + ) + + if key.shape[-2] != add_key.shape[-2] or value.shape[-2] != add_value.shape[-2]: + raise ValueError( + "Additional key/value tokens must match the self attention key/value tokens in number " + f"of heads, got {key.shape[-2]=} != {add_key.shape[-2]=}, and " + f"{value.shape[-2]=} != {add_value.shape[-2]=}." + ) + + if key.shape[-1] != add_key.shape[-1] or value.shape[-1] != add_value.shape[-1]: + raise ValueError( + "Additional key/value tokens must match the self attention key/value tokens in head " + f"dim, got {key.shape[-1]=} != {add_key.shape[-1]=}, and " + f"{value.shape[-1]=} != {add_value.shape[-1]=}." + ) + + if not supports_gqa_mqa and ( + query.shape[-2] != add_key.shape[-2] or query.shape[-2] != add_value.shape[-2] + ): + raise ValueError( + f"This operation does not support GQA/MQA, therefore number of heads in Q, K, and V " + f"must match, got {query.shape[-2]=}, {key.shape[-2]=}, {value.shape[-2]=}." + ) + + if supports_gqa_mqa: + if ( + key.shape[-2] != value.shape[-2] + or key.shape[-2] != add_key.shape[-2] + or key.shape[-2] != add_value.shape[-2] + ): + raise ValueError( + "Key and value, original and additional, must always have the same number of heads, got " + f"{key.shape[-2]=}, {value.shape[-2]=}, {add_key.shape[-2]=}, {add_value.shape[-2]=}." + ) + + heads_q = query.shape[-2] + heads_kv = key.shape[-2] + + if heads_q < heads_kv or heads_q % heads_kv != 0: + raise ValueError( + "Key/value heads must evenly divide query heads, got " + f"{heads_q=}, {heads_kv=}." + ) + + +def check_input_size_arg(na_dim: int, input_size: Any) -> DimensionType: + assert na_dim > 0 and na_dim < 4 + if ( + isinstance(input_size, Sequence) + and len(input_size) == na_dim + and all(isinstance(x, int) and x > 1 for x in input_size) + ): + return tuple(x for x in input_size) + + if isinstance(input_size, int) and input_size > 1: + return create_dim_from_int(na_dim, value=input_size) + + raise ValueError( + "Invalid value for `input_size`; expected an integer or iterable of integers, all >= 2, " + f"got {type(input_size)=}, {input_size=}." + ) + + +def check_kernel_size_arg(na_dim: int, kernel_size: Any) -> DimensionType: + assert na_dim > 0 and na_dim < 4 + if ( + isinstance(kernel_size, Sequence) + and len(kernel_size) == na_dim + and all(isinstance(x, int) and x > 1 for x in kernel_size) + ): + return tuple(x for x in kernel_size) + + if isinstance(kernel_size, int) and kernel_size > 1: + return create_dim_from_int(na_dim, value=kernel_size) + + raise ValueError( + "Invalid value for `kernel_size`; expected an integer or iterable of integers, all >= 2, " + f"got {type(kernel_size)=}, {kernel_size=}." + ) + + +def check_stride_arg(na_dim: int, stride: Any) -> DimensionType: + assert na_dim > 0 and na_dim < 4 + if stride is None: + return create_dim_from_int(na_dim, value=1) + + if ( + isinstance(stride, Sequence) + and len(stride) == na_dim + and all(isinstance(x, int) and x > 0 for x in stride) + ): + return tuple(x for x in stride) + + if isinstance(stride, int) and stride > 0: + return create_dim_from_int(na_dim, value=stride) + + raise ValueError( + "Invalid value for `stride`; expected an integer or tuple of positive integers, " + f"got {type(stride)=}, {stride=}." + ) + + +def check_dilation_arg(na_dim: int, dilation: Any) -> DimensionType: + assert na_dim > 0 and na_dim < 4 + if dilation is None: + return create_dim_from_int(na_dim, value=1) + + if ( + isinstance(dilation, Sequence) + and len(dilation) == na_dim + and all(isinstance(x, int) and x > 0 for x in dilation) + ): + return tuple(x for x in dilation) + + if isinstance(dilation, int) and dilation > 0: + return create_dim_from_int(na_dim, value=dilation) + + raise ValueError( + "Invalid value for `dilation`; expected an integer or tuple of positive integers, " + f"got {type(dilation)=}, {dilation=}." + ) + + +def check_causal_arg(na_dim: int, is_causal: Any) -> CausalArgType: + assert na_dim > 0 and na_dim < 4 + + if is_causal is None: + return create_causal_arg_from_bool(na_dim, value=False) + + if ( + isinstance(is_causal, Sequence) + and len(is_causal) == na_dim + and all(isinstance(c, bool) for c in is_causal) + ): + return tuple(c for c in is_causal) + + if isinstance(is_causal, bool): + return create_causal_arg_from_bool(na_dim, value=is_causal) + + raise ValueError( + "Invalid value for `is_causal`; expected a boolean or tuple of booleans, " + f"got {type(is_causal)=}, {is_causal=}." + ) + + +def check_all_args( + na_dim: int, kernel_size: Any, stride: Any, dilation: Any, is_causal: Any +) -> Tuple[DimensionType, DimensionType, DimensionType, CausalArgType]: + kernel_size_out, stride_out, dilation_out, is_causal_out = ( + check_kernel_size_arg(na_dim, kernel_size), + check_stride_arg(na_dim, stride), + check_dilation_arg(na_dim, dilation), + check_causal_arg(na_dim, is_causal), + ) + + return kernel_size_out, stride_out, dilation_out, is_causal_out + + +def check_args_against_input( + input_tensor: Tensor, + kernel_size: DimensionType, + stride: DimensionType, + dilation: DimensionType, + is_causal: CausalArgType, +): + assert input_tensor.dim() in [4, 5, 6] + na_dim = input_tensor.dim() - 3 + input_size = input_tensor.shape[1 : 1 + na_dim] + + if any(k * d > x for x, k, d in zip(input_size, kernel_size, dilation)): + raise ValueError( + "The product of kernel size and dilation cannot be larger than input size " + f"along any dimension, got {input_size=} ({input_tensor.shape=}), " + f"{kernel_size=}, {dilation=}." + ) + + if any(s > k for k, s in zip(kernel_size, stride)): + raise ValueError( + "Stride cannot be larger than kernel size along any dimension, got " + f"{kernel_size=}, {stride=}." + ) + + +def is_self_attention( + input_tensor: Tensor, + kernel_size: DimensionType, + is_causal: CausalArgType, + has_additional_attention: bool, +): + assert input_tensor.dim() in [4, 5, 6] + na_dim = input_tensor.dim() - 3 + input_size = input_tensor.shape[1 : 1 + na_dim] + + # Special case: 1-D causal with full window is equivalent to standard 1-D causal + # as long as there isn't any additional context (non causal) + if na_dim == 1 and not has_additional_attention: + return kernel_size[0] == input_size[0] + + return all(k == x and not c for x, k, c in zip(input_size, kernel_size, is_causal)) + + +def check_tile_shape( + tile_shape: Any, +) -> DimensionType: + if ( + isinstance(tile_shape, Sequence) + and len(tile_shape) <= 3 + and all(isinstance(x, int) for x in tile_shape) + ): + return tuple(x for x in tile_shape) + + raise ValueError( + f"Unsupported value for tile shape; expected an iterable of at most 3 integers, " + f"got {type(tile_shape)=}, {tile_shape}." + ) + + +def check_kernel_schedule(kernel_schedule: Any) -> Optional[KernelSchedule]: + if kernel_schedule is None: + return None + + if isinstance(kernel_schedule, KernelSchedule): + return kernel_schedule + + if kernel_schedule == "non": + return KernelSchedule.NonPersistent + elif kernel_schedule == "coop": + return KernelSchedule.WarpSpecializedCooperative + elif kernel_schedule == "pp": + return KernelSchedule.WarpSpecializedPingpong + + raise ValueError( + f"Kernel schedule {kernel_schedule} is invalid; choices are: " + "`non` (non-persistent), `coop` (warp-specialized cooperative), and " + "`pp` (warp-specialized ping-ponging)." + ) + + +# Varlen FMHA Checks + + +def varlen_tensor_checks( + query: Tensor, + key: Tensor, + value: Tensor, + seqlens_Q: Optional[Tensor] = None, + seqlens_KV: Optional[Tensor] = None, + cumulative_seqlen_Q: Optional[Tensor] = None, + cumulative_seqlen_KV: Optional[Tensor] = None, + max_seqlen_Q: Optional[int] = None, + max_seqlen_KV: Optional[int] = None, +) -> Union[ + Tuple[NoneType, NoneType, int, int], + Tuple[Tensor, Tensor, int, int], +]: + if query.shape[0] != key.shape[0] or query.shape[0] != value.shape[0]: + raise ValueError( + "Q, K, and V must match in batch size, got " + f"{query.shape[0]=}, {key.shape[0]=}, {value.shape[0]=}." + ) + + if all( + x is None + for x in [ + seqlens_Q, + seqlens_KV, + cumulative_seqlen_Q, + cumulative_seqlen_KV, + ] + ) and all( + x is None or x == 0 + for x in [ + max_seqlen_Q, + max_seqlen_KV, + ] + ): + # Not varlen + return None, None, 0, 0 + + if seqlens_Q is not None or seqlens_KV is not None: + # Generate cumulative_seqlen_{Q,KV}, max_seqlen_{Q,KV}, total_seqlen_{Q,KV} + # based on user input + return generate_varlen_parameters( + query=query, + key=key, + value=value, + seqlens_Q=seqlens_Q, + seqlens_KV=seqlens_KV, + ) + + # Validate user-input cumulative_seqlen_{Q,KV}, max_seqlen_{Q,KV}, total_seqlen_{Q,KV} + if any( + x is None + for x in [ + cumulative_seqlen_Q, + cumulative_seqlen_KV, + max_seqlen_Q, + max_seqlen_KV, + ] + ): + raise ValueError( + "Variable length Attention requires all of " + "cumulative_seqlen_{Q,KV} and max_seqlen_{Q,KV} to be set." + ) + + if query.shape[0] != 1: + raise ValueError( + "Variable length Attention only supports sequence-packed memory layout " + f"(batch = 1), got {query.shape[0]=}." + ) + + assert cumulative_seqlen_Q is not None + assert cumulative_seqlen_KV is not None + assert max_seqlen_Q is not None + assert max_seqlen_KV is not None + + if not isinstance(max_seqlen_Q, int) or not isinstance(max_seqlen_KV, int): + raise ValueError( + "max_seqlen_Q and max_seqlen_KV must be ints, got " + f"{type(max_seqlen_Q)=}, {type(max_seqlen_KV)=}, {max_seqlen_Q=}, {max_seqlen_KV=}." + ) + + total_seqlen_Q = query.shape[1] + total_seqlen_KV = key.shape[1] + if max_seqlen_Q > total_seqlen_Q: + raise ValueError( + "Maximum sequence length cannot exceed total, got " + f"{max_seqlen_Q=}, {total_seqlen_Q=}." + ) + + if max_seqlen_KV > total_seqlen_KV: + raise ValueError( + "Maximum sequence length cannot exceed total, got " + f"{max_seqlen_KV=}, {total_seqlen_KV=}." + ) + + # NOTE: this check introduces recompiles + if not is_torch_compiling(): + if (max_seqlen_Q == 0) != (max_seqlen_KV == 0): + raise ValueError( + "max_seqlen_Q and max_seqlen_KV must both be zero or both be non-zero, got " + f"{max_seqlen_Q=}, {max_seqlen_KV=}." + ) + + if max_seqlen_Q < 0 or max_seqlen_KV < 0: + raise ValueError( + "Maximum sequence length cannot be negative, got " + f"{max_seqlen_Q=}, {max_seqlen_KV=}." + ) + + if not isinstance(cumulative_seqlen_Q, Tensor) or not isinstance( + cumulative_seqlen_KV, Tensor + ): + raise ValueError( + "cumulative_seqlen_Q and cumulative_seqlen_KV must both be tensors." + ) + + if ( + cumulative_seqlen_Q.device != query.device + or cumulative_seqlen_KV.device != query.device + ): + raise ValueError( + "cumulative_seqlen_Q and cumulative_seqlen_KV must be on the same device as QKV, but " + f"{cumulative_seqlen_Q.device=}, {cumulative_seqlen_KV.device=}, {query.device=}." + ) + + if ( + cumulative_seqlen_Q.dtype != torch.int32 + or cumulative_seqlen_KV.dtype != torch.int32 + ): + raise ValueError( + "cumulative_seqlen_Q and cumulative_seqlen_KV must both be torch.int32 tensors, got " + f"{cumulative_seqlen_Q.dtype=}, {cumulative_seqlen_KV.dtype=}." + ) + + if cumulative_seqlen_Q.dim() != 1 or cumulative_seqlen_KV.dim() != 1: + raise ValueError( + "cumulative_seqlen_Q and cumulative_seqlen_KV must both be 1-D tensors, got " + f"{cumulative_seqlen_Q.dim()=}, {cumulative_seqlen_KV.dim()=}." + ) + + if cumulative_seqlen_Q.shape[0] != cumulative_seqlen_KV.shape[0]: + raise ValueError( + "cumulative_seqlen_Q and cumulative_seqlen_KV must match in size, got " + f"{cumulative_seqlen_Q.shape=}, {cumulative_seqlen_KV.shape=}." + ) + + if cumulative_seqlen_Q.shape[0] < 2: + raise ValueError( + "cumulative_seqlen_Q and cumulative_seqlen_KV must contain at least 2 elements, got " + f"{cumulative_seqlen_Q.shape=}, {cumulative_seqlen_KV.shape=}." + ) + + return ( + cumulative_seqlen_Q, + cumulative_seqlen_KV, + max_seqlen_Q, + max_seqlen_KV, + ) diff --git a/natten/torch-ext/natten/utils/device.py b/natten/torch-ext/natten/utils/device.py new file mode 100644 index 00000000..7997ab47 --- /dev/null +++ b/natten/torch-ext/natten/utils/device.py @@ -0,0 +1,50 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +from typing import Optional + +import torch + + +def is_cuda(device: torch.device) -> bool: + return torch.cuda.is_available() and torch.version.cuda and device.type == "cuda" # type: ignore + + +def is_rocm(device: torch.device) -> bool: + return torch.cuda.is_available() and torch.version.hip and device.type == "cuda" # type: ignore + + +def is_cpu(device: torch.device) -> bool: + return device.type == "cpu" + + +def get_device_cc(device: Optional[torch.device] = None) -> int: + if ( + torch.cuda.is_available() + and torch.version.cuda + and (device is None or is_cuda(device)) + ): + major, minor = torch.cuda.get_device_capability(device) + return major * 10 + minor + + return 0 diff --git a/natten/torch-ext/natten/utils/dtype.py b/natten/torch-ext/natten/utils/dtype.py new file mode 100644 index 00000000..c65d9e9e --- /dev/null +++ b/natten/torch-ext/natten/utils/dtype.py @@ -0,0 +1,36 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +import torch + + +def is_full(dtype: torch.dtype) -> bool: + return dtype == torch.float32 + + +def is_half(dtype: torch.dtype) -> bool: + return dtype in [torch.float16, torch.bfloat16] + + +def is_fp8(dtype: torch.dtype) -> bool: + return dtype in [torch.float8_e5m2, torch.float8_e4m3fn] diff --git a/natten/torch-ext/natten/utils/environment.py b/natten/torch-ext/natten/utils/environment.py new file mode 100644 index 00000000..0581aacb --- /dev/null +++ b/natten/torch-ext/natten/utils/environment.py @@ -0,0 +1,79 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +import os + +import torch + +from ..utils.device import get_device_cc + + +def parse_env_flag(env_var: str, default: bool) -> bool: + default_str = "1" if default else "0" + out_str = os.getenv(env_var, default_str) + if out_str.strip() == "": + return default + if out_str == "0": + return False + if out_str == "1": + return True + return default + + +def parse_env_int(env_var: str, default: int) -> int: + out_str = os.getenv(env_var, str(default)) + if out_str.strip() == "": + return default + try: + return int(out_str) + except ValueError: + return default + + +def parse_env_str(env_var: str, default: str) -> str: + return os.getenv(env_var, str(default)) + + +_IS_CUDA_AVAILABLE = torch.cuda.is_available() + +_TORCH_VERSION = [int(x) for x in torch.__version__.split(".")[:2]] + +_IS_TORCH_COMPILE_SUPPORTED = _TORCH_VERSION >= [2, 6] and get_device_cc() >= 70 + +# Guard registering libnatten APIs as torch ops with environment variables +# In case any unusual bugs from torch compile come up again +# Also restrict to torch 2.8 and later +# https://github.com/pytorch/pytorch/issues/137979#issuecomment-3614956989 +DISABLE_TORCH_OPS = _TORCH_VERSION < [2, 8] or parse_env_flag( + "NATTEN_DISABLE_TORCH_OPS", False +) + + +# Controls all regions guarded against torch compile +# Logs, and certain assertions cause graph breaks. +def is_torch_compiling() -> bool: + try: + return torch.compiler.is_compiling() + except: + # Assume too old to support torch compile + return False diff --git a/natten/torch-ext/natten/utils/log.py b/natten/torch-ext/natten/utils/log.py new file mode 100644 index 00000000..138c3ac8 --- /dev/null +++ b/natten/torch-ext/natten/utils/log.py @@ -0,0 +1,134 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +import enum +import logging +import os +import sys + +from ..utils.environment import is_torch_compiling, parse_env_str + +log_format = "| %(asctime)s | [[ %(name)s ]] [ %(levelname)s ]: %(message)s" + + +class LogLevel(enum.Enum): + Default = 0 + Debug = 1 + Info = 2 + Warnings = 3 + Errors = 4 + Critical = 5 + + +def _get_log_level() -> LogLevel: + log_level = parse_env_str("NATTEN_LOG_LEVEL", "").lower() + + if log_level == "debug": + return LogLevel.Debug + elif log_level == "info": + return LogLevel.Info + elif log_level == "warning": + return LogLevel.Warnings + elif log_level == "error": + return LogLevel.Errors + elif log_level == "critical": + return LogLevel.Critical + + return LogLevel.Default + + +_map_log_level = { + LogLevel.Default: logging.INFO, + LogLevel.Debug: logging.DEBUG, + LogLevel.Info: logging.INFO, + LogLevel.Warnings: logging.WARNING, + LogLevel.Errors: logging.ERROR, + LogLevel.Critical: logging.CRITICAL, +} + + +# Tests will stream into stderr instead of stdout +# It can be set to either stderr, stdout or any writeable file. +# Otherwise logging will be disabled. +def _get_log_pipe(): + log_pipe = parse_env_str("NATTEN_LOG_PIPE", "stdout") + + # Skip checking /dev/null writablity + if log_pipe == "/dev/null": + return None + + if log_pipe.lower() == "stderr": + return sys.stderr + + if log_pipe.lower() == "stdout": + return sys.stdout + + # Treat as file path; validate writability + if os.path.isfile(log_pipe) and os.access(log_pipe, os.W_OK): + return log_pipe + + try: + open(log_pipe, "a").close() + return log_pipe + except OSError: + pass + + return None + + +class NattenLogger: + def __init__(self, name: str): + self.logger = logging.getLogger(name) + self.log_level = _map_log_level[_get_log_level()] + self.logger.setLevel(self.log_level) + self.formatter = logging.Formatter(log_format) + log_pipe = _get_log_pipe() + if log_pipe in [sys.stderr, sys.stdout]: + self.handler = logging.StreamHandler(log_pipe) + elif isinstance(log_pipe, str): + self.handler = logging.FileHandler(log_pipe) + else: + # Invalid / null + self.handler = logging.NullHandler() # type: ignore[assignment] + self.handler.setLevel(self.log_level) + self.handler.setFormatter(self.formatter) + self.logger.addHandler(self.handler) + + def is_safe_to_log(self) -> bool: + return not is_torch_compiling() + + def info(self, *args, **kwargs): + if self.is_safe_to_log(): + self.logger.info(*args, **kwargs) + + def debug(self, *args, **kwargs): + if self.is_safe_to_log(): + self.logger.debug(*args, **kwargs) + + def warning(self, *args, **kwargs): + if self.is_safe_to_log(): + self.logger.warning(*args, **kwargs) + + +def get_logger(name) -> NattenLogger: + return NattenLogger(name) diff --git a/natten/torch-ext/natten/utils/tensor.py b/natten/torch-ext/natten/utils/tensor.py new file mode 100644 index 00000000..32d95acf --- /dev/null +++ b/natten/torch-ext/natten/utils/tensor.py @@ -0,0 +1,113 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +from typing import Optional + +import torch +from torch import Size, Tensor + + +def _get_expected_attn_shape(input_tensor: Tensor, attention_dim: int) -> Size: + shape = [x for x in input_tensor.shape[:-1]] + [attention_dim] + return Size(shape) + + +def make_attn_tensor_from_input(input_tensor: Tensor, attention_dim: int) -> Tensor: + return torch.empty( + _get_expected_attn_shape(input_tensor, attention_dim), + device=input_tensor.device, + dtype=input_tensor.dtype, + requires_grad=input_tensor.requires_grad, + ) + + +def check_additional_keys( + input_tensor: Tensor, additional_keys: Optional[Tensor] +) -> int: + if additional_keys is None: + return 0 + + if additional_keys.dim() != 4: + raise ValueError( + "Additional tokens have to be shaped as a rank-4 tensor; " + f"got {additional_keys.dim()}." + ) + batch_size, heads, tokens, dim = additional_keys.shape + expected_batch_size = input_tensor.shape[0] + expected_heads = input_tensor.shape[1] + expected_dim = input_tensor.shape[-1] + if ( + batch_size != expected_batch_size + or expected_heads != heads + or expected_dim != dim + ): + raise ValueError( + "Shape mismatch between input tensor and additional tokens; " + "they must match in batch size, heads, and dim per head. " + f"Got {input_tensor.shape=}, {additional_keys.shape=}." + ) + return tokens + + +def check_additional_values( + attn_tensor: Tensor, + additional_values: Optional[Tensor], + value: Tensor, + expected_attn_weights: int, +) -> int: + if additional_values is None and attn_tensor.shape[-1] == expected_attn_weights: + return 0 + if additional_values is None: + raise ValueError( + f"Expected {expected_attn_weights} attention weights per token, " + f"got {attn_tensor.shape[-1]=}." + ) + + if additional_values.dim() != 4: + raise ValueError( + "Additional tokens have to be shaped as a rank-4 tensor; " + f"got {additional_values.dim()}." + ) + + if additional_values.shape[-1] != value.shape[-1]: + raise ValueError( + "Additional value tokens must match the dimension of the " + f"rest of the tokens, got {additional_values.shape[-1]=} != " + f"{value.shape[-1]=}." + ) + + batch_size, heads, tokens, dim = additional_values.shape + if tokens + expected_attn_weights != attn_tensor.shape[-1]: + raise ValueError( + f"Expected {expected_attn_weights + tokens} attention weights per token, " + f"got {attn_tensor.shape[-1]=}." + ) + expected_batch_size = attn_tensor.shape[0] + expected_heads = attn_tensor.shape[1] + if batch_size != expected_batch_size or expected_heads != heads: + raise ValueError( + "Shape mismatch between attention tensor and additional tokens; " + "they must match in batch size and heads. " + f"Got {attn_tensor.shape=}, {additional_values.shape=}." + ) + return tokens diff --git a/natten/torch-ext/natten/utils/testing.py b/natten/torch-ext/natten/utils/testing.py new file mode 100644 index 00000000..8ed9de5d --- /dev/null +++ b/natten/torch-ext/natten/utils/testing.py @@ -0,0 +1,149 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +import torch + +from .._environment import _IS_CUDA_AVAILABLE, _RUN_EXTENDED_TESTS, HAS_LIBNATTEN +from ..backends.flex import _FLEX_COMPILE_SUPPORTED, _FLEX_SUPPORTED +from ..utils.device import get_device_cc, is_cuda + + +def skip_if_libnatten_is_not_supported(): + def decorator(f): + def wrapper(self, *args, **kwargs): + if not _IS_CUDA_AVAILABLE: + self.skipTest("CUDA is not available.") + elif not HAS_LIBNATTEN: + self.skipTest("Libnatten is not available.") + else: + return f(self, *args, **kwargs) + + return wrapper + + return decorator + + +def skip_if_cuda_is_not_supported(): + def decorator(f): + def wrapper(self, *args, **kwargs): + if not _IS_CUDA_AVAILABLE: + self.skipTest("CUDA is not available.") + else: + return f(self, *args, **kwargs) + + return wrapper + + return decorator + + +def skip_if_flex_is_not_supported(): + def decorator(f): + def wrapper(self, *args, **kwargs): + if not _FLEX_SUPPORTED or get_device_cc() < 70: + self.skipTest("Flex backend is not supported.") + else: + return f(self, *args, **kwargs) + + return wrapper + + return decorator + + +def skip_if_flex_compile_is_not_supported(): + def decorator(f): + def wrapper(self, *args, **kwargs): + if not _FLEX_COMPILE_SUPPORTED: + self.skipTest("Flex (compiled) backend is not supported.") + else: + return f(self, *args, **kwargs) + + return wrapper + + return decorator + + +def skip_if_not_running_extended_tests(): + def decorator(f): + def wrapper(self, *args, **kwargs): + if not _RUN_EXTENDED_TESTS: + self.skipTest("Skipping extended test cases.") + else: + return f(self, *args, **kwargs) + + return wrapper + + return decorator + + +def skip_if_hopper_kernels_not_supported(): + def decorator(f): + def wrapper(self, *args, **kwargs): + if get_device_cc() != 90: + self.skipTest("Hopper kernels are only supported on SM90.") + else: + return f(self, *args, **kwargs) + + return wrapper + + return decorator + + +def skip_if_blackwell_kernels_not_supported(): + def decorator(f): + def wrapper(self, *args, **kwargs): + if get_device_cc() not in [100, 103]: + self.skipTest( + "Blackwell kernels are only supported on SM100 and SM103." + ) + else: + return f(self, *args, **kwargs) + + return wrapper + + return decorator + + +def supports_float16(device: torch.device) -> bool: + if is_cuda(device): + device_cc = get_device_cc(device) + + if device_cc < 50: + return False + + return True + + # TODO: + return True + + +def supports_bfloat16(device: torch.device) -> bool: + if is_cuda(device): + device_cc = get_device_cc(device) + + if device_cc < 80: + return False + + return True + + # TODO: + return False diff --git a/natten/torch-ext/natten/utils/tuples.py b/natten/torch-ext/natten/utils/tuples.py new file mode 100644 index 00000000..d1f16de7 --- /dev/null +++ b/natten/torch-ext/natten/utils/tuples.py @@ -0,0 +1,51 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +from ..types import CausalArgType, DimensionType + + +def ceil_div_int(x: int, y: int) -> int: + return (x + y - 1) // y + + +def ceil_div_tuple(X: tuple, Y: tuple) -> tuple: + assert len(X) == len(Y) + return tuple(ceil_div_int(x, y) for x, y in zip(X, Y)) + + +def mul_tuple(X: tuple, Y: tuple) -> tuple: + assert len(X) == len(Y) + return tuple(x * y for x, y in zip(X, Y)) + + +def sub_tuple(X: tuple, Y: tuple) -> tuple: + assert len(X) == len(Y) + return tuple(x - y for x, y in zip(X, Y)) + + +def create_dim_from_int(na_dim: int, value: int) -> DimensionType: + return tuple(value for _ in range(na_dim)) # type: ignore + + +def create_causal_arg_from_bool(na_dim: int, value: bool) -> CausalArgType: + return tuple(value for _ in range(na_dim)) # type: ignore diff --git a/natten/torch-ext/natten/utils/varlen.py b/natten/torch-ext/natten/utils/varlen.py new file mode 100644 index 00000000..9d4e35d9 --- /dev/null +++ b/natten/torch-ext/natten/utils/varlen.py @@ -0,0 +1,135 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +from typing import Optional, Tuple, Union + +import torch # noqa: F401 +from torch import Tensor + +from ..types import NoneType +from ..utils.environment import is_torch_compiling + + +def generate_varlen_parameters( + query: Tensor, + key: Tensor, + value: Tensor, + seqlens_Q: Optional[Tensor] = None, + seqlens_KV: Optional[Tensor] = None, +) -> Union[ + Tuple[NoneType, NoneType, int, int], + Tuple[Tensor, Tensor, int, int], +]: + # NOTE: max_seqlen_{Q,KV} require a device-host sync, since they're expected to be ints (with + # which we launch the varlen kernel) and not device tensors. + # .item() introduces control flow and breaks the graph. + # It is also inefficient to repeat this per-op, and mostly there for convenience. + # generate_varlen_parameters should ideally always be called by the user ahead of model + # forward / backward. + if is_torch_compiling(): + raise RuntimeError( + "Running 'generate_varlen_parameters' in a torch-compiled region is disallowed as it " + "results in graph breaks. Please consider calling ahead of time and pass " + "'cumulative_seqlen_{Q,KV}' and 'max_seqlen_{Q,KV}' instead of 'seqlens_{Q,KV}' to " + "'attention'. " + ) + + if query.shape[0] != key.shape[0] or query.shape[0] != value.shape[0]: + raise ValueError( + "Q, K, and V must match in batch size, got " + f"{query.shape[0]=}, {key.shape[0]=}, {value.shape[0]=}." + ) + + if (seqlens_Q is None) ^ (seqlens_KV is None): + raise ValueError( + "Variable length Attention requires both of seqlens_Q and seqlens_KV to be set, got " + f"{seqlens_Q=}, {seqlens_KV=}." + ) + + if seqlens_Q is None and seqlens_KV is None: + # Not varlen + return None, None, 0, 0 + + assert seqlens_Q is not None + assert seqlens_KV is not None + + if not isinstance(seqlens_Q, Tensor) or not isinstance(seqlens_KV, Tensor): + raise ValueError("seqlens_Q and seqlens_KV must both be tensors.") + + if seqlens_Q.device != query.device or seqlens_KV.device != query.device: + raise ValueError( + "seqlens_Q and seqlens_KV must be on the same device as QKV, but " + f"{seqlens_Q.device=}, {seqlens_KV.device=}, {query.device=}." + ) + + if seqlens_Q.dtype != torch.int32 or seqlens_KV.dtype != torch.int32: + raise ValueError( + "seqlens_Q and seqlens_KV must both be torch.int32 tensors, got " + f"{seqlens_Q.dtype=}, {seqlens_KV.dtype=}." + ) + + if seqlens_Q.dim() != 1 or seqlens_KV.dim() != 1: + raise ValueError( + "seqlens_Q and seqlens_KV must both be 1-D tensors, got " + f"{seqlens_Q.dim()=}, {seqlens_KV.dim()=}." + ) + + if seqlens_Q.shape[0] != seqlens_KV.shape[0]: + raise ValueError( + "seqlens_Q and seqlens_KV must match in size, got " + f"{seqlens_Q.shape=}, {seqlens_KV.shape=}." + ) + + if seqlens_Q.shape[0] < 1: + raise ValueError( + "seqlens_Q and seqlens_KV must contain at least one element, got " + f"{seqlens_Q.shape=}, {seqlens_KV.shape=}." + ) + + if query.shape[0] != 1: + raise ValueError( + "Variable length attention only supports sequence-packed memory layout " + f"(batch = 1), got {query.shape[0]=}." + ) + + assert seqlens_Q.dim() == seqlens_KV.dim() == 1 + assert seqlens_Q.shape[0] == seqlens_KV.shape[0] >= 1 + assert seqlens_Q.dtype == seqlens_KV.dtype == torch.int32 + + max_seqlen_Q = seqlens_Q.max().item() # type: ignore + max_seqlen_KV = seqlens_KV.max().item() # type: ignore + + # NOTE: we have to prepend with 0 manually :( + z = torch.tensor([0], dtype=torch.int32, device=seqlens_Q.device) + cumulative_seqlen_Q = torch.cat([z, seqlens_Q.cumsum(0).to(torch.int32)], dim=0) + cumulative_seqlen_KV = torch.cat([z, seqlens_KV.cumsum(0).to(torch.int32)], dim=0) + + assert isinstance(max_seqlen_Q, int) + assert isinstance(max_seqlen_KV, int) + + return ( + cumulative_seqlen_Q, + cumulative_seqlen_KV, + max_seqlen_Q, + max_seqlen_KV, + ) diff --git a/natten/torch-ext/natten/version.py b/natten/torch-ext/natten/version.py new file mode 100644 index 00000000..846fd135 --- /dev/null +++ b/natten/torch-ext/natten/version.py @@ -0,0 +1,24 @@ +################################################################################################# +# Copyright (c) 2022 - 2026 Ali Hassani. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################################# + +__version__ = "0.21.6" diff --git a/natten/torch-ext/torch_binding.cpp b/natten/torch-ext/torch_binding.cpp new file mode 100644 index 00000000..fd768c5f --- /dev/null +++ b/natten/torch-ext/torch_binding.cpp @@ -0,0 +1,871 @@ +/*************************************************************************************************** + * Copyright (c) 2022 - 2026 Ali Hassani. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + **************************************************************************************************/ +// kernel-builder port: replaces upstream csrc/natten.cpp (pybind11) with +// TORCH_LIBRARY registrations against the Python limited API. The op schemas +// are out-variant and use `int[]` for NATTEN's dimension tuples (including +// 0/1 encoded causal masks); the adapters below convert to the +// std::tuple signatures of the NATTEN C++ entry points. + +#include +#include + +#include +#include +#include + +#include "natten/blackwell_fmha.h" +#include "natten/blackwell_fna.h" +#include "natten/compute_delta.h" +#include "natten/fmha.h" +#include "natten/fna.h" +#include "natten/hopper_fmha.h" +#include "natten/hopper_fna.h" +#include "natten/reference.h" +#include "natten/token_permute.h" + +#include "registration.h" + +namespace { + +using IntVec = std::vector; +using OptTensor = std::optional; + +template +auto to_int32_tuple(const IntVec &v) { + TORCH_CHECK( + v.size() == N, + "NATTEN: expected an int list of length ", + N, + ", got ", + v.size()); + if constexpr (N == 1) { + return std::make_tuple(static_cast(v[0])); + } else if constexpr (N == 2) { + return std::make_tuple( + static_cast(v[0]), static_cast(v[1])); + } else { + static_assert(N == 3); + return std::make_tuple( + static_cast(v[0]), + static_cast(v[1]), + static_cast(v[2])); + } +} + +template +auto to_bool_tuple(const IntVec &v) { + TORCH_CHECK( + v.size() == N, + "NATTEN: expected a causal mask list of length ", + N, + ", got ", + v.size()); + if constexpr (N == 1) { + return std::make_tuple(v[0] != 0); + } else if constexpr (N == 2) { + return std::make_tuple(v[0] != 0, v[1] != 0); + } else { + static_assert(N == 3); + return std::make_tuple(v[0] != 0, v[1] != 0, v[2] != 0); + } +} + +// Legacy (CUTLASS 2.X) fused neighborhood attention. + +#define DEFINE_FNA_ADAPTERS(DIM) \ + void na##DIM##d_forward( \ + at::Tensor &out, \ + const at::Tensor &query, \ + const at::Tensor &key, \ + const at::Tensor &value, \ + at::Tensor &logsumexp, \ + IntVec kernel_size, \ + IntVec stride, \ + IntVec dilation, \ + IntVec is_causal, \ + double scale, \ + IntVec q_tile_shape, \ + IntVec kv_tile_shape) { \ + natten::na##DIM##d_forward( \ + out, \ + query, \ + key, \ + value, \ + at::optional(logsumexp), \ + to_int32_tuple(kernel_size), \ + to_int32_tuple(stride), \ + to_int32_tuple(dilation), \ + to_bool_tuple(is_causal), \ + static_cast(scale), \ + to_int32_tuple(q_tile_shape), \ + to_int32_tuple(kv_tile_shape)); \ + } \ + \ + void na##DIM##d_backward( \ + at::Tensor &grad_query, \ + at::Tensor &grad_key, \ + at::Tensor &grad_value, \ + const at::Tensor &query, \ + const at::Tensor &key, \ + const at::Tensor &value, \ + const at::Tensor &out, \ + const at::Tensor &grad_out, \ + const at::Tensor &logsumexp, \ + IntVec kernel_size, \ + IntVec stride, \ + IntVec dilation, \ + IntVec is_causal, \ + double scale, \ + IntVec q_tile_shape, \ + IntVec kv_tile_shape, \ + IntVec num_kv_splits, \ + bool compute_delta_with_torch) { \ + natten::na##DIM##d_backward( \ + grad_query, \ + grad_key, \ + grad_value, \ + query, \ + key, \ + value, \ + out, \ + grad_out, \ + logsumexp, \ + to_int32_tuple(kernel_size), \ + to_int32_tuple(stride), \ + to_int32_tuple(dilation), \ + to_bool_tuple(is_causal), \ + static_cast(scale), \ + to_int32_tuple(q_tile_shape), \ + to_int32_tuple(kv_tile_shape), \ + to_int32_tuple(num_kv_splits), \ + compute_delta_with_torch); \ + } + +DEFINE_FNA_ADAPTERS(1) +DEFINE_FNA_ADAPTERS(2) +DEFINE_FNA_ADAPTERS(3) + +#undef DEFINE_FNA_ADAPTERS + +// Hopper (SM90) fused neighborhood attention. + +#define DEFINE_HOPPER_FNA_ADAPTERS(DIM) \ + void hopper_na##DIM##d_forward( \ + at::Tensor &out, \ + const at::Tensor &query, \ + const at::Tensor &key, \ + const at::Tensor &value, \ + at::Tensor &logsumexp, \ + IntVec kernel_size, \ + IntVec stride, \ + IntVec dilation, \ + IntVec is_causal, \ + double scale, \ + IntVec q_shape, \ + IntVec kv_shape, \ + IntVec qkv_shape, \ + IntVec q_tile_shape, \ + IntVec kv_tile_shape, \ + int64_t kernel_type) { \ + natten::hopper_na##DIM##d_forward( \ + out, \ + query, \ + key, \ + value, \ + at::optional(logsumexp), \ + to_int32_tuple(kernel_size), \ + to_int32_tuple(stride), \ + to_int32_tuple(dilation), \ + to_bool_tuple(is_causal), \ + static_cast(scale), \ + to_int32_tuple(q_shape), \ + to_int32_tuple(kv_shape), \ + to_int32_tuple(qkv_shape), \ + to_int32_tuple(q_tile_shape), \ + to_int32_tuple(kv_tile_shape), \ + static_cast(kernel_type)); \ + } \ + \ + void hopper_na##DIM##d_backward( \ + at::Tensor &grad_query, \ + at::Tensor &grad_key, \ + at::Tensor &grad_value, \ + const at::Tensor &query, \ + const at::Tensor &key, \ + const at::Tensor &value, \ + const at::Tensor &out, \ + const at::Tensor &grad_out, \ + const at::Tensor &logsumexp, \ + IntVec kernel_size, \ + IntVec stride, \ + IntVec dilation, \ + IntVec is_causal, \ + double scale, \ + IntVec q_shape, \ + IntVec kv_shape, \ + IntVec qkv_shape, \ + IntVec q_tile_shape, \ + IntVec kv_tile_shape) { \ + natten::hopper_na##DIM##d_backward( \ + grad_query, \ + grad_key, \ + grad_value, \ + query, \ + key, \ + value, \ + out, \ + grad_out, \ + logsumexp, \ + to_int32_tuple(kernel_size), \ + to_int32_tuple(stride), \ + to_int32_tuple(dilation), \ + to_bool_tuple(is_causal), \ + static_cast(scale), \ + to_int32_tuple(q_shape), \ + to_int32_tuple(kv_shape), \ + to_int32_tuple(qkv_shape), \ + to_int32_tuple(q_tile_shape), \ + to_int32_tuple(kv_tile_shape)); \ + } + +DEFINE_HOPPER_FNA_ADAPTERS(1) +DEFINE_HOPPER_FNA_ADAPTERS(2) +DEFINE_HOPPER_FNA_ADAPTERS(3) + +#undef DEFINE_HOPPER_FNA_ADAPTERS + +// Blackwell (SM100) fused neighborhood attention. + +#define DEFINE_BLACKWELL_FNA_ADAPTERS(DIM) \ + void blackwell_na##DIM##d_forward( \ + at::Tensor &out, \ + const at::Tensor &query, \ + const at::Tensor &key, \ + const at::Tensor &value, \ + at::Tensor &logsumexp, \ + IntVec kernel_size, \ + IntVec stride, \ + IntVec dilation, \ + IntVec is_causal, \ + double scale, \ + IntVec q_shape, \ + IntVec kv_shape, \ + IntVec qkv_shape, \ + IntVec q_tile_shape, \ + IntVec kv_tile_shape, \ + bool run_persistent) { \ + natten::blackwell_na##DIM##d_forward( \ + out, \ + query, \ + key, \ + value, \ + at::optional(logsumexp), \ + to_int32_tuple(kernel_size), \ + to_int32_tuple(stride), \ + to_int32_tuple(dilation), \ + to_bool_tuple(is_causal), \ + static_cast(scale), \ + to_int32_tuple(q_shape), \ + to_int32_tuple(kv_shape), \ + to_int32_tuple(qkv_shape), \ + to_int32_tuple(q_tile_shape), \ + to_int32_tuple(kv_tile_shape), \ + run_persistent); \ + } \ + \ + void blackwell_na##DIM##d_backward( \ + at::Tensor &grad_query, \ + at::Tensor &grad_key, \ + at::Tensor &grad_value, \ + const at::Tensor &query, \ + const at::Tensor &key, \ + const at::Tensor &value, \ + const at::Tensor &out, \ + const at::Tensor &grad_out, \ + const at::Tensor &logsumexp, \ + IntVec kernel_size, \ + IntVec stride, \ + IntVec dilation, \ + IntVec is_causal, \ + double scale, \ + IntVec q_shape, \ + IntVec kv_shape, \ + IntVec qkv_shape, \ + IntVec q_tile_shape, \ + IntVec kv_tile_shape) { \ + natten::blackwell_na##DIM##d_backward( \ + grad_query, \ + grad_key, \ + grad_value, \ + query, \ + key, \ + value, \ + out, \ + grad_out, \ + logsumexp, \ + to_int32_tuple(kernel_size), \ + to_int32_tuple(stride), \ + to_int32_tuple(dilation), \ + to_bool_tuple(is_causal), \ + static_cast(scale), \ + to_int32_tuple(q_shape), \ + to_int32_tuple(kv_shape), \ + to_int32_tuple(qkv_shape), \ + to_int32_tuple(q_tile_shape), \ + to_int32_tuple(kv_tile_shape)); \ + } + +DEFINE_BLACKWELL_FNA_ADAPTERS(1) +DEFINE_BLACKWELL_FNA_ADAPTERS(2) +DEFINE_BLACKWELL_FNA_ADAPTERS(3) + +#undef DEFINE_BLACKWELL_FNA_ADAPTERS + +// Reference kernels. + +#define DEFINE_REFERENCE_FNA_ADAPTERS(DIM) \ + void reference_na##DIM##d_forward( \ + at::Tensor &out, \ + const at::Tensor &query, \ + const at::Tensor &key, \ + const at::Tensor &value, \ + at::Tensor &logsumexp, \ + IntVec kernel_size, \ + IntVec stride, \ + IntVec dilation, \ + IntVec is_causal, \ + double scale, \ + IntVec qkv_shape, \ + int64_t num_extra_kv) { \ + natten::reference_na##DIM##d_forward( \ + out, \ + query, \ + key, \ + value, \ + logsumexp, \ + to_int32_tuple(kernel_size), \ + to_int32_tuple(stride), \ + to_int32_tuple(dilation), \ + to_bool_tuple(is_causal), \ + static_cast(scale), \ + to_int32_tuple(qkv_shape), \ + static_cast(num_extra_kv)); \ + } \ + \ + void reference_na##DIM##d_backward( \ + at::Tensor &grad_query, \ + at::Tensor &grad_key, \ + at::Tensor &grad_value, \ + const at::Tensor &query, \ + const at::Tensor &key, \ + const at::Tensor &value, \ + const at::Tensor &out, \ + const at::Tensor &grad_out, \ + const at::Tensor &logsumexp, \ + IntVec kernel_size, \ + IntVec stride, \ + IntVec dilation, \ + IntVec is_causal, \ + double scale, \ + IntVec qkv_shape, \ + int64_t num_extra_kv) { \ + natten::reference_na##DIM##d_backward( \ + grad_query, \ + grad_key, \ + grad_value, \ + query, \ + key, \ + value, \ + out, \ + grad_out, \ + logsumexp, \ + to_int32_tuple(kernel_size), \ + to_int32_tuple(stride), \ + to_int32_tuple(dilation), \ + to_bool_tuple(is_causal), \ + static_cast(scale), \ + to_int32_tuple(qkv_shape), \ + static_cast(num_extra_kv)); \ + } + +DEFINE_REFERENCE_FNA_ADAPTERS(1) +DEFINE_REFERENCE_FNA_ADAPTERS(2) +DEFINE_REFERENCE_FNA_ADAPTERS(3) + +#undef DEFINE_REFERENCE_FNA_ADAPTERS + +// Token permute kernels. + +#define DEFINE_TOKEN_PERMUTE_ADAPTERS(DIM) \ + void token_permute_##DIM##d( \ + at::Tensor &out, \ + const at::Tensor &in, \ + IntVec tile_shape, \ + IntVec dilation, \ + bool flip_tiled_dims) { \ + natten::token_permute_##DIM##d( \ + out, \ + in, \ + to_int32_tuple(tile_shape), \ + to_int32_tuple(dilation), \ + flip_tiled_dims); \ + } \ + \ + void token_unpermute_##DIM##d( \ + at::Tensor &out, \ + const at::Tensor &in, \ + IntVec tile_shape, \ + IntVec dilation, \ + bool flip_tiled_dims) { \ + natten::token_unpermute_##DIM##d( \ + out, \ + in, \ + to_int32_tuple(tile_shape), \ + to_int32_tuple(dilation), \ + flip_tiled_dims); \ + } + +DEFINE_TOKEN_PERMUTE_ADAPTERS(1) +DEFINE_TOKEN_PERMUTE_ADAPTERS(2) +DEFINE_TOKEN_PERMUTE_ADAPTERS(3) + +#undef DEFINE_TOKEN_PERMUTE_ADAPTERS + +// FMHA (CUTLASS 2.X). + +void fmha_forward( + at::Tensor &out, + const at::Tensor &query, + const at::Tensor &key, + const at::Tensor &value, + at::Tensor &logsumexp, + bool is_causal, + double scale, + int64_t q_tile_size, + int64_t kv_tile_size, + const OptTensor &cumulative_seqlen_Q, + const OptTensor &cumulative_seqlen_KV, + int64_t max_seqlen_Q, + int64_t max_seqlen_KV) { + natten::fmha_forward( + out, + query, + key, + value, + at::optional(logsumexp), + is_causal, + static_cast(scale), + static_cast(q_tile_size), + static_cast(kv_tile_size), + cumulative_seqlen_Q, + cumulative_seqlen_KV, + static_cast(max_seqlen_Q), + static_cast(max_seqlen_KV)); +} + +void fmha_backward( + at::Tensor &grad_query, + at::Tensor &grad_key, + at::Tensor &grad_value, + const at::Tensor &query, + const at::Tensor &key, + const at::Tensor &value, + const at::Tensor &out, + const at::Tensor &grad_out, + const at::Tensor &logsumexp, + bool is_causal, + double scale, + int64_t q_tile_size, + int64_t kv_tile_size, + int64_t num_kv_splits, + bool compute_delta_with_torch, + const OptTensor &cumulative_seqlen_Q, + const OptTensor &cumulative_seqlen_KV, + int64_t max_seqlen_Q, + int64_t max_seqlen_KV) { + natten::fmha_backward( + grad_query, + grad_key, + grad_value, + query, + key, + value, + out, + grad_out, + logsumexp, + is_causal, + static_cast(scale), + static_cast(q_tile_size), + static_cast(kv_tile_size), + static_cast(num_kv_splits), + compute_delta_with_torch, + cumulative_seqlen_Q, + cumulative_seqlen_KV, + static_cast(max_seqlen_Q), + static_cast(max_seqlen_KV)); +} + +// Hopper (SM90) FMHA. + +void hopper_fmha_forward( + at::Tensor &out, + const at::Tensor &query, + const at::Tensor &key, + const at::Tensor &value, + at::Tensor &logsumexp, + bool is_causal, + double scale, + int64_t q_tile_size, + int64_t kv_tile_size, + int64_t kernel_type, + const OptTensor &cumulative_seqlen_Q, + const OptTensor &cumulative_seqlen_KV, + int64_t max_seqlen_Q, + int64_t max_seqlen_KV) { + natten::hopper_fmha_forward( + out, + query, + key, + value, + at::optional(logsumexp), + is_causal, + static_cast(scale), + static_cast(q_tile_size), + static_cast(kv_tile_size), + static_cast(kernel_type), + cumulative_seqlen_Q, + cumulative_seqlen_KV, + static_cast(max_seqlen_Q), + static_cast(max_seqlen_KV)); +} + +void hopper_fmha_backward( + at::Tensor &grad_query, + at::Tensor &grad_key, + at::Tensor &grad_value, + const at::Tensor &query, + const at::Tensor &key, + const at::Tensor &value, + const at::Tensor &out, + const at::Tensor &grad_out, + const at::Tensor &logsumexp, + bool is_causal, + double scale, + int64_t q_tile_size, + int64_t kv_tile_size, + const OptTensor &cumulative_seqlen_Q, + const OptTensor &cumulative_seqlen_KV, + int64_t max_seqlen_Q, + int64_t max_seqlen_KV) { + natten::hopper_fmha_backward( + grad_query, + grad_key, + grad_value, + query, + key, + value, + out, + grad_out, + logsumexp, + is_causal, + static_cast(scale), + static_cast(q_tile_size), + static_cast(kv_tile_size), + cumulative_seqlen_Q, + cumulative_seqlen_KV, + static_cast(max_seqlen_Q), + static_cast(max_seqlen_KV)); +} + +// Blackwell (SM100) FMHA. + +void blackwell_fmha_forward( + at::Tensor &out, + const at::Tensor &query, + const at::Tensor &key, + const at::Tensor &value, + at::Tensor &logsumexp, + bool is_causal, + double scale, + int64_t q_tile_size, + int64_t kv_tile_size, + bool run_persistent, + const OptTensor &cumulative_seqlen_Q, + const OptTensor &cumulative_seqlen_KV, + int64_t max_seqlen_Q, + int64_t max_seqlen_KV) { + natten::blackwell_fmha_forward( + out, + query, + key, + value, + at::optional(logsumexp), + is_causal, + static_cast(scale), + static_cast(q_tile_size), + static_cast(kv_tile_size), + run_persistent, + cumulative_seqlen_Q, + cumulative_seqlen_KV, + static_cast(max_seqlen_Q), + static_cast(max_seqlen_KV)); +} + +void blackwell_fmha_backward( + at::Tensor &grad_query, + at::Tensor &grad_key, + at::Tensor &grad_value, + const at::Tensor &query, + const at::Tensor &key, + const at::Tensor &value, + const at::Tensor &out, + const at::Tensor &grad_out, + const at::Tensor &logsumexp, + bool is_causal, + double scale, + int64_t q_tile_size, + int64_t kv_tile_size, + const OptTensor &cumulative_seqlen_Q, + const OptTensor &cumulative_seqlen_KV, + int64_t max_seqlen_Q, + int64_t max_seqlen_KV, + bool deterministic) { + natten::blackwell_fmha_backward( + grad_query, + grad_key, + grad_value, + query, + key, + value, + out, + grad_out, + logsumexp, + is_causal, + static_cast(scale), + static_cast(q_tile_size), + static_cast(kv_tile_size), + cumulative_seqlen_Q, + cumulative_seqlen_KV, + static_cast(max_seqlen_Q), + static_cast(max_seqlen_KV), + deterministic); +} + +// Misc. + +void compute_delta( + const at::Tensor &out, + const at::Tensor &d_out, + at::Tensor &delta) { + natten::compute_delta(out, d_out, delta); +} + +} // namespace + +TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { + // Legacy (CUTLASS 2.X) FNA +#define REGISTER_FNA_OPS(DIM) \ + ops.def( \ + "na" #DIM "d_forward(Tensor! out, Tensor query, Tensor key, " \ + "Tensor value, Tensor! logsumexp, int[] kernel_size, int[] stride, " \ + "int[] dilation, int[] is_causal, float scale, " \ + "int[] q_tile_shape, int[] kv_tile_shape) -> ()"); \ + ops.impl("na" #DIM "d_forward", torch::kCUDA, &na##DIM##d_forward); \ + ops.def( \ + "na" #DIM "d_backward(Tensor! grad_query, Tensor! grad_key, " \ + "Tensor! grad_value, Tensor query, Tensor key, Tensor value, " \ + "Tensor out, Tensor grad_out, Tensor logsumexp, " \ + "int[] kernel_size, int[] stride, int[] dilation, int[] is_causal, " \ + "float scale, int[] q_tile_shape, int[] kv_tile_shape, " \ + "int[] num_kv_splits, bool compute_delta_with_torch) -> ()"); \ + ops.impl("na" #DIM "d_backward", torch::kCUDA, &na##DIM##d_backward); + + REGISTER_FNA_OPS(1) + REGISTER_FNA_OPS(2) + REGISTER_FNA_OPS(3) +#undef REGISTER_FNA_OPS + + // Hopper (SM90) FNA +#define REGISTER_HOPPER_FNA_OPS(DIM) \ + ops.def( \ + "hopper_na" #DIM "d_forward(Tensor! out, Tensor query, Tensor key, " \ + "Tensor value, Tensor! logsumexp, int[] kernel_size, int[] stride, " \ + "int[] dilation, int[] is_causal, float scale, int[] q_shape, " \ + "int[] kv_shape, int[] qkv_shape, int[] q_tile_shape, " \ + "int[] kv_tile_shape, int kernel_type) -> ()"); \ + ops.impl( \ + "hopper_na" #DIM "d_forward", \ + torch::kCUDA, \ + &hopper_na##DIM##d_forward); \ + ops.def( \ + "hopper_na" #DIM "d_backward(Tensor! grad_query, Tensor! grad_key, " \ + "Tensor! grad_value, Tensor query, Tensor key, Tensor value, " \ + "Tensor out, Tensor grad_out, Tensor logsumexp, " \ + "int[] kernel_size, int[] stride, int[] dilation, int[] is_causal, " \ + "float scale, int[] q_shape, int[] kv_shape, int[] qkv_shape, " \ + "int[] q_tile_shape, int[] kv_tile_shape) -> ()"); \ + ops.impl( \ + "hopper_na" #DIM "d_backward", \ + torch::kCUDA, \ + &hopper_na##DIM##d_backward); + + REGISTER_HOPPER_FNA_OPS(1) + REGISTER_HOPPER_FNA_OPS(2) + REGISTER_HOPPER_FNA_OPS(3) +#undef REGISTER_HOPPER_FNA_OPS + + // Blackwell (SM100) FNA +#define REGISTER_BLACKWELL_FNA_OPS(DIM) \ + ops.def( \ + "blackwell_na" #DIM "d_forward(Tensor! out, Tensor query, " \ + "Tensor key, Tensor value, Tensor! logsumexp, int[] kernel_size, " \ + "int[] stride, int[] dilation, int[] is_causal, float scale, " \ + "int[] q_shape, int[] kv_shape, int[] qkv_shape, " \ + "int[] q_tile_shape, int[] kv_tile_shape, " \ + "bool run_persistent) -> ()"); \ + ops.impl( \ + "blackwell_na" #DIM "d_forward", \ + torch::kCUDA, \ + &blackwell_na##DIM##d_forward); \ + ops.def( \ + "blackwell_na" #DIM "d_backward(Tensor! grad_query, " \ + "Tensor! grad_key, Tensor! grad_value, Tensor query, Tensor key, " \ + "Tensor value, Tensor out, Tensor grad_out, Tensor logsumexp, " \ + "int[] kernel_size, int[] stride, int[] dilation, int[] is_causal, " \ + "float scale, int[] q_shape, int[] kv_shape, int[] qkv_shape, " \ + "int[] q_tile_shape, int[] kv_tile_shape) -> ()"); \ + ops.impl( \ + "blackwell_na" #DIM "d_backward", \ + torch::kCUDA, \ + &blackwell_na##DIM##d_backward); + + REGISTER_BLACKWELL_FNA_OPS(1) + REGISTER_BLACKWELL_FNA_OPS(2) + REGISTER_BLACKWELL_FNA_OPS(3) +#undef REGISTER_BLACKWELL_FNA_OPS + + // Reference kernels +#define REGISTER_REFERENCE_FNA_OPS(DIM) \ + ops.def( \ + "reference_na" #DIM "d_forward(Tensor! out, Tensor query, " \ + "Tensor key, Tensor value, Tensor! logsumexp, int[] kernel_size, " \ + "int[] stride, int[] dilation, int[] is_causal, float scale, " \ + "int[] qkv_shape, int num_extra_kv) -> ()"); \ + ops.impl( \ + "reference_na" #DIM "d_forward", \ + torch::kCUDA, \ + &reference_na##DIM##d_forward); \ + ops.def( \ + "reference_na" #DIM "d_backward(Tensor! grad_query, " \ + "Tensor! grad_key, Tensor! grad_value, Tensor query, Tensor key, " \ + "Tensor value, Tensor out, Tensor grad_out, Tensor logsumexp, " \ + "int[] kernel_size, int[] stride, int[] dilation, int[] is_causal, " \ + "float scale, int[] qkv_shape, int num_extra_kv) -> ()"); \ + ops.impl( \ + "reference_na" #DIM "d_backward", \ + torch::kCUDA, \ + &reference_na##DIM##d_backward); + + REGISTER_REFERENCE_FNA_OPS(1) + REGISTER_REFERENCE_FNA_OPS(2) + REGISTER_REFERENCE_FNA_OPS(3) +#undef REGISTER_REFERENCE_FNA_OPS + + // Token permute +#define REGISTER_TOKEN_PERMUTE_OPS(DIM) \ + ops.def( \ + "token_permute_" #DIM "d(Tensor! out, Tensor input, " \ + "int[] tile_shape, int[] dilation, bool flip_tiled_dims) -> ()"); \ + ops.impl( \ + "token_permute_" #DIM "d", torch::kCUDA, &token_permute_##DIM##d); \ + ops.def( \ + "token_unpermute_" #DIM "d(Tensor! out, Tensor input, " \ + "int[] tile_shape, int[] dilation, bool flip_tiled_dims) -> ()"); \ + ops.impl( \ + "token_unpermute_" #DIM "d", \ + torch::kCUDA, \ + &token_unpermute_##DIM##d); + + REGISTER_TOKEN_PERMUTE_OPS(1) + REGISTER_TOKEN_PERMUTE_OPS(2) + REGISTER_TOKEN_PERMUTE_OPS(3) +#undef REGISTER_TOKEN_PERMUTE_OPS + + // FMHA + ops.def( + "fmha_forward(Tensor! out, Tensor query, Tensor key, Tensor value, " + "Tensor! logsumexp, bool is_causal, float scale, int q_tile_size, " + "int kv_tile_size, Tensor? cumulative_seqlen_Q, " + "Tensor? cumulative_seqlen_KV, int max_seqlen_Q, " + "int max_seqlen_KV) -> ()"); + ops.impl("fmha_forward", torch::kCUDA, &fmha_forward); + ops.def( + "fmha_backward(Tensor! grad_query, Tensor! grad_key, " + "Tensor! grad_value, Tensor query, Tensor key, Tensor value, " + "Tensor out, Tensor grad_out, Tensor logsumexp, bool is_causal, " + "float scale, int q_tile_size, int kv_tile_size, int num_kv_splits, " + "bool compute_delta_with_torch, Tensor? cumulative_seqlen_Q, " + "Tensor? cumulative_seqlen_KV, int max_seqlen_Q, " + "int max_seqlen_KV) -> ()"); + ops.impl("fmha_backward", torch::kCUDA, &fmha_backward); + + // Hopper FMHA + ops.def( + "hopper_fmha_forward(Tensor! out, Tensor query, Tensor key, " + "Tensor value, Tensor! logsumexp, bool is_causal, float scale, " + "int q_tile_size, int kv_tile_size, int kernel_type, " + "Tensor? cumulative_seqlen_Q, Tensor? cumulative_seqlen_KV, " + "int max_seqlen_Q, int max_seqlen_KV) -> ()"); + ops.impl("hopper_fmha_forward", torch::kCUDA, &hopper_fmha_forward); + ops.def( + "hopper_fmha_backward(Tensor! grad_query, Tensor! grad_key, " + "Tensor! grad_value, Tensor query, Tensor key, Tensor value, " + "Tensor out, Tensor grad_out, Tensor logsumexp, bool is_causal, " + "float scale, int q_tile_size, int kv_tile_size, " + "Tensor? cumulative_seqlen_Q, Tensor? cumulative_seqlen_KV, " + "int max_seqlen_Q, int max_seqlen_KV) -> ()"); + ops.impl("hopper_fmha_backward", torch::kCUDA, &hopper_fmha_backward); + + // Blackwell FMHA + ops.def( + "blackwell_fmha_forward(Tensor! out, Tensor query, Tensor key, " + "Tensor value, Tensor! logsumexp, bool is_causal, float scale, " + "int q_tile_size, int kv_tile_size, bool run_persistent, " + "Tensor? cumulative_seqlen_Q, Tensor? cumulative_seqlen_KV, " + "int max_seqlen_Q, int max_seqlen_KV) -> ()"); + ops.impl("blackwell_fmha_forward", torch::kCUDA, &blackwell_fmha_forward); + ops.def( + "blackwell_fmha_backward(Tensor! grad_query, Tensor! grad_key, " + "Tensor! grad_value, Tensor query, Tensor key, Tensor value, " + "Tensor out, Tensor grad_out, Tensor logsumexp, bool is_causal, " + "float scale, int q_tile_size, int kv_tile_size, " + "Tensor? cumulative_seqlen_Q, Tensor? cumulative_seqlen_KV, " + "int max_seqlen_Q, int max_seqlen_KV, bool deterministic) -> ()"); + ops.impl("blackwell_fmha_backward", torch::kCUDA, &blackwell_fmha_backward); + + // Misc + ops.def("compute_delta(Tensor out, Tensor d_out, Tensor! delta) -> ()"); + ops.impl("compute_delta", torch::kCUDA, &compute_delta); +} + +REGISTER_EXTENSION(TORCH_EXTENSION_NAME) From 755ffeec6a9e4c8ecc6e737aea9f458094948f3a Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Thu, 16 Jul 2026 16:15:08 +0530 Subject: [PATCH 2/8] natten: add to KERNEL_SOURCE_MAPPING Co-Authored-By: Claude Fable 5 --- scripts/check_kernel_freshness.py | 1 + 1 file changed, 1 insertion(+) diff --git a/scripts/check_kernel_freshness.py b/scripts/check_kernel_freshness.py index 977a3f1a..3acb8f48 100644 --- a/scripts/check_kernel_freshness.py +++ b/scripts/check_kernel_freshness.py @@ -29,6 +29,7 @@ "megablocks": "https://github.com/databricks/megablocks", "mra": "", "msa": "https://github.com/MiniMax-AI/MSA", + "natten": "https://github.com/SHI-Labs/NATTEN", "paged-attention": "", "punica-sgmv": "https://github.com/predibase/lorax", "quantization-bitsandbytes": "https://github.com/bitsandbytes-foundation/bitsandbytes", From 545ffce5ca67ba2af694471f17b1c9cabeed0c51 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Thu, 16 Jul 2026 16:27:15 +0530 Subject: [PATCH 3/8] natten: pin kernel-builder past cutlass_4_5 hash fix CI failed on all variants with a fixed-output hash mismatch fetching the CUTLASS 4.5.2 source: nix-builder's recorded hash predates whatever the v4.5.2 tag resolves to now, and natten is the first cutlass_4_5 consumer. Pin the kernel-builder input to huggingface/kernels@ce92bd77 (branch fix-cutlass-4-5-hash), which corrects the hash; re-lock to main once the upstream PR merges. Co-Authored-By: Claude Fable 5 --- natten/flake.lock | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/natten/flake.lock b/natten/flake.lock index d7fdf232..a58cc3b9 100644 --- a/natten/flake.lock +++ b/natten/flake.lock @@ -41,11 +41,11 @@ "rust-overlay": "rust-overlay" }, "locked": { - "lastModified": 1783417137, - "narHash": "sha256-/iWe/m5ZACdb/DcXTrvXmbHS2QrzgGcGa7hPKi460tE=", + "lastModified": 1784199293, + "narHash": "sha256-hYPTttTha+J8nJCX9PsrPkz8S34fJ1c4AI17po6cNQg=", "owner": "huggingface", "repo": "kernels", - "rev": "81580bb92577f2f7228661ca2a221fc052375709", + "rev": "ce92bd77c807ce3a2f4b85d8bab69caf027b51d0", "type": "github" }, "original": { @@ -56,17 +56,17 @@ }, "nixpkgs": { "locked": { - "lastModified": 1776927958, - "narHash": "sha256-XOzEtft7E0P6TgQViLUOQeGHlEYiQ0+FY24BPEksj6s=", + "lastModified": 1783284758, + "narHash": "sha256-tiQ8/qi8I45OOaBBYlVbXoAVkeQzvvTQOv5I45rMw5o=", "owner": "NixOS", "repo": "nixpkgs", - "rev": "fec2c46cca5bf9767486a290abae51200b656d69", + "rev": "ec1a11210589d294f0ac99d3290a27e6c73dfa1d", "type": "github" }, "original": { "owner": "NixOS", - "ref": "nixos-unstable-small", "repo": "nixpkgs", + "rev": "ec1a11210589d294f0ac99d3290a27e6c73dfa1d", "type": "github" } }, @@ -83,11 +83,11 @@ ] }, "locked": { - "lastModified": 1776914043, - "narHash": "sha256-qug5r56yW1qOsjSI99l3Jm15JNT9CvS2otkXNRNtrPI=", + "lastModified": 1783320166, + "narHash": "sha256-l7C/OsjcnWDOk2K3ssj+SBduwL67LashjBqis9+t468=", "owner": "oxalica", "repo": "rust-overlay", - "rev": "2d35c4358d7de3a0e606a6e8b27925d981c01cc3", + "rev": "20ee15370c9256669d66968b89ee20a4b0a4e673", "type": "github" }, "original": { From 39160cf398fdf251faf03ede7fa8d56dfa29701e Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Thu, 16 Jul 2026 16:45:44 +0530 Subject: [PATCH 4/8] natten: replace torch:: dtype constants with at:: equivalents Upstream NATTEN's helpers.h and the autogen templates use torch::kFloat16 & co., which come from the torch C++ frontend headers that upstream pulled in transitively via torch/extension.h. With that include stripped for the limited API, the autogen TUs only see ATen and fail to compile (https://github.com/huggingface/kernels-community/actions/runs/29492671685). Switch to the identical at:: constants (kFloat32->kFloat, kFloat16->kHalf, kBFloat16->kBFloat16, zeros->zeros) in helpers.h and the autogen scripts, and regenerate. Co-Authored-By: Claude Fable 5 --- .../cuda/blackwell_fmha/interface.h | 4 +- .../cuda/blackwell_fmha_bwd/interface.h | 4 +- .../cuda/blackwell_fna/dispatch_dtype.h | 12 +- .../cuda/blackwell_fna_bwd/dispatch_dtype.h | 12 +- .../natten_autogen/cuda/fmha/dispatch_dtype.h | 36 +++--- .../natten_autogen/cuda/fna/dispatch_dtype.h | 108 +++++++++--------- .../cuda/hopper_fmha/interface.h | 4 +- .../cuda/hopper_fmha_bwd/interface.h | 4 +- .../cuda/hopper_fna/dispatch_dtype.h | 12 +- .../cuda/hopper_fna_bwd/dispatch_dtype.h | 12 +- .../cuda/reference/dispatch_dtype.h | 36 +++--- .../src/cuda/blackwell_fmha_bwd/source_0.cu | 4 +- .../src/cuda/blackwell_fmha_bwd/source_1.cu | 4 +- .../src/cuda/blackwell_fmha_bwd/source_2.cu | 2 +- .../src/cuda/blackwell_fmha_bwd/source_3.cu | 2 +- natten/csrc/include/natten/helpers.h | 16 +-- natten/scripts/autogen_blackwell_fmha.py | 4 +- natten/scripts/autogen_blackwell_fmha_bwd.py | 6 +- natten/scripts/autogen_blackwell_fna.py | 4 +- natten/scripts/autogen_blackwell_fna_bwd.py | 4 +- natten/scripts/autogen_fmha.py | 6 +- natten/scripts/autogen_fna.py | 6 +- natten/scripts/autogen_hopper_fmha.py | 4 +- natten/scripts/autogen_hopper_fmha_bwd.py | 4 +- natten/scripts/autogen_hopper_fna.py | 4 +- natten/scripts/autogen_hopper_fna_bwd.py | 4 +- natten/scripts/autogen_reference_fna.py | 6 +- 27 files changed, 162 insertions(+), 162 deletions(-) diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/interface.h b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/interface.h index 17357193..5cc68419 100644 --- a/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/interface.h +++ b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/interface.h @@ -21,10 +21,10 @@ namespace cuda { namespace fmha_blackwell { #define DISPATCH_BLACKWELL_FMHA_FORWARD(dtype, dim, q_tile_size, kv_tile_size, persistent, ...) \ [&] { \ - if (dtype == torch::kFloat16) { \ + if (dtype == at::kHalf) { \ DISPATCH_BLACKWELL_FMHA_FORWARD_float16(dim, q_tile_size, kv_tile_size, persistent, __VA_ARGS__); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ DISPATCH_BLACKWELL_FMHA_FORWARD_bfloat16(dim, q_tile_size, kv_tile_size, persistent, __VA_ARGS__); \ } \ else if (dtype == c10::ScalarType::Float8_e4m3fn) { \ diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/interface.h b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/interface.h index 103277bc..38e91ef0 100644 --- a/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/interface.h +++ b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/interface.h @@ -21,10 +21,10 @@ namespace cuda { namespace fmha_blackwell { #define DISPATCH_BLACKWELL_FMHA_BACKWARD(dtype, dim, q_tile_size, kv_tile_size, ...) \ [&] { \ - if (dtype == torch::kFloat16) { \ + if (dtype == at::kHalf) { \ DISPATCH_BLACKWELL_FMHA_BACKWARD_float16(dim, q_tile_size, kv_tile_size, __VA_ARGS__); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ DISPATCH_BLACKWELL_FMHA_BACKWARD_bfloat16(dim, q_tile_size, kv_tile_size, __VA_ARGS__); \ } \ else { \ diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_dtype.h b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_dtype.h index 2e3b6392..ca4a9afb 100644 --- a/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_dtype.h +++ b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_dtype.h @@ -21,10 +21,10 @@ namespace cuda { namespace fna_blackwell { #define DISPATCH_BLACKWELL_FNA_FORWARD_1D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ [&] { \ - if (dtype == torch::kFloat16) { \ + if (dtype == at::kHalf) { \ DISPATCH_BLACKWELL_FNA_FORWARD_1D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ DISPATCH_BLACKWELL_FNA_FORWARD_1D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ } \ else if (dtype == c10::ScalarType::Float8_e4m3fn) { \ @@ -40,10 +40,10 @@ namespace fna_blackwell { #define DISPATCH_BLACKWELL_FNA_FORWARD_2D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ [&] { \ - if (dtype == torch::kFloat16) { \ + if (dtype == at::kHalf) { \ DISPATCH_BLACKWELL_FNA_FORWARD_2D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ DISPATCH_BLACKWELL_FNA_FORWARD_2D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ } \ else if (dtype == c10::ScalarType::Float8_e4m3fn) { \ @@ -59,10 +59,10 @@ namespace fna_blackwell { #define DISPATCH_BLACKWELL_FNA_FORWARD_3D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, persistent, ...) \ [&] { \ - if (dtype == torch::kFloat16) { \ + if (dtype == at::kHalf) { \ DISPATCH_BLACKWELL_FNA_FORWARD_3D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ DISPATCH_BLACKWELL_FNA_FORWARD_3D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, persistent, __VA_ARGS__); \ } \ else if (dtype == c10::ScalarType::Float8_e4m3fn) { \ diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_dtype.h b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_dtype.h index 5ca7f4e1..b691164c 100644 --- a/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_dtype.h +++ b/natten/csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_dtype.h @@ -21,10 +21,10 @@ namespace cuda { namespace fna_blackwell { #define DISPATCH_BLACKWELL_FNA_BACKWARD_1D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, ...) \ [&] { \ - if (dtype == torch::kFloat16) { \ + if (dtype == at::kHalf) { \ DISPATCH_BLACKWELL_FNA_BACKWARD_1D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ DISPATCH_BLACKWELL_FNA_BACKWARD_1D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ } \ else { \ @@ -34,10 +34,10 @@ namespace fna_blackwell { #define DISPATCH_BLACKWELL_FNA_BACKWARD_2D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, ...) \ [&] { \ - if (dtype == torch::kFloat16) { \ + if (dtype == at::kHalf) { \ DISPATCH_BLACKWELL_FNA_BACKWARD_2D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ DISPATCH_BLACKWELL_FNA_BACKWARD_2D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ } \ else { \ @@ -47,10 +47,10 @@ namespace fna_blackwell { #define DISPATCH_BLACKWELL_FNA_BACKWARD_3D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, ...) \ [&] { \ - if (dtype == torch::kFloat16) { \ + if (dtype == at::kHalf) { \ DISPATCH_BLACKWELL_FNA_BACKWARD_3D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ DISPATCH_BLACKWELL_FNA_BACKWARD_3D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ } \ else { \ diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/fmha/dispatch_dtype.h b/natten/csrc/autogen/include/natten_autogen/cuda/fmha/dispatch_dtype.h index c11a5dc1..ed663184 100644 --- a/natten/csrc/autogen/include/natten_autogen/cuda/fmha/dispatch_dtype.h +++ b/natten/csrc/autogen/include/natten_autogen/cuda/fmha/dispatch_dtype.h @@ -13,10 +13,10 @@ namespace cuda { namespace fmha { #define DISPATCH_FMHA_FORWARD_SM50(dtype, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ fmha_sm50_float32(cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ fmha_sm50_float16(cb); \ } \ else { \ @@ -26,10 +26,10 @@ namespace fmha { #define DISPATCH_FMHA_FORWARD_SM70(dtype, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ fmha_sm70_float32(cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ fmha_sm70_float16(cb); \ } \ else { \ @@ -39,10 +39,10 @@ namespace fmha { #define DISPATCH_FMHA_FORWARD_SM75(dtype, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ fmha_sm75_float32(cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ fmha_sm75_float16(cb); \ } \ else { \ @@ -52,13 +52,13 @@ namespace fmha { #define DISPATCH_FMHA_FORWARD_SM80(dtype, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ fmha_sm80_float32(cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ fmha_sm80_float16(cb); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ fmha_sm80_bfloat16(cb); \ } \ else { \ @@ -68,10 +68,10 @@ namespace fmha { #define DISPATCH_FMHA_BACKWARD_SM50(dtype, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ fmha_backward_sm50_float32(cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ fmha_backward_sm50_float16(cb); \ } \ else { \ @@ -81,10 +81,10 @@ namespace fmha { #define DISPATCH_FMHA_BACKWARD_SM70(dtype, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ fmha_backward_sm70_float32(cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ fmha_backward_sm70_float16(cb); \ } \ else { \ @@ -94,10 +94,10 @@ namespace fmha { #define DISPATCH_FMHA_BACKWARD_SM75(dtype, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ fmha_backward_sm75_float32(cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ fmha_backward_sm75_float16(cb); \ } \ else { \ @@ -107,13 +107,13 @@ namespace fmha { #define DISPATCH_FMHA_BACKWARD_SM80(dtype, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ fmha_backward_sm80_float32(cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ fmha_backward_sm80_float16(cb); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ fmha_backward_sm80_bfloat16(cb); \ } \ else { \ diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/fna/dispatch_dtype.h b/natten/csrc/autogen/include/natten_autogen/cuda/fna/dispatch_dtype.h index cbbb324c..bd342c0a 100644 --- a/natten/csrc/autogen/include/natten_autogen/cuda/fna/dispatch_dtype.h +++ b/natten/csrc/autogen/include/natten_autogen/cuda/fna/dispatch_dtype.h @@ -14,10 +14,10 @@ namespace cuda { namespace fna { #define DISPATCH_FNA_FORWARD_1D_SM50(dtype, is_causal, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_FNA_FORWARD_1D_SM50_float32(is_causal, cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_FNA_FORWARD_1D_SM50_float16(is_causal, cb); \ } \ else { \ @@ -27,10 +27,10 @@ namespace fna { #define DISPATCH_FNA_FORWARD_1D_SM70(dtype, is_causal, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_FNA_FORWARD_1D_SM70_float32(is_causal, cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_FNA_FORWARD_1D_SM70_float16(is_causal, cb); \ } \ else { \ @@ -40,10 +40,10 @@ namespace fna { #define DISPATCH_FNA_FORWARD_1D_SM75(dtype, is_causal, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_FNA_FORWARD_1D_SM75_float32(is_causal, cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_FNA_FORWARD_1D_SM75_float16(is_causal, cb); \ } \ else { \ @@ -53,13 +53,13 @@ namespace fna { #define DISPATCH_FNA_FORWARD_1D_SM80(dtype, is_causal, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_FNA_FORWARD_1D_SM80_float32(is_causal, cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_FNA_FORWARD_1D_SM80_float16(is_causal, cb); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ DISPATCH_FNA_FORWARD_1D_SM80_bfloat16(is_causal, cb); \ } \ else { \ @@ -69,10 +69,10 @@ namespace fna { #define DISPATCH_FNA_FORWARD_2D_SM50(dtype, is_causal, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_FNA_FORWARD_2D_SM50_float32(is_causal, cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_FNA_FORWARD_2D_SM50_float16(is_causal, cb); \ } \ else { \ @@ -82,10 +82,10 @@ namespace fna { #define DISPATCH_FNA_FORWARD_2D_SM70(dtype, is_causal, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_FNA_FORWARD_2D_SM70_float32(is_causal, cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_FNA_FORWARD_2D_SM70_float16(is_causal, cb); \ } \ else { \ @@ -95,10 +95,10 @@ namespace fna { #define DISPATCH_FNA_FORWARD_2D_SM75(dtype, is_causal, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_FNA_FORWARD_2D_SM75_float32(is_causal, cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_FNA_FORWARD_2D_SM75_float16(is_causal, cb); \ } \ else { \ @@ -108,13 +108,13 @@ namespace fna { #define DISPATCH_FNA_FORWARD_2D_SM80(dtype, is_causal, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_FNA_FORWARD_2D_SM80_float32(is_causal, cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_FNA_FORWARD_2D_SM80_float16(is_causal, cb); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ DISPATCH_FNA_FORWARD_2D_SM80_bfloat16(is_causal, cb); \ } \ else { \ @@ -124,10 +124,10 @@ namespace fna { #define DISPATCH_FNA_FORWARD_3D_SM50(dtype, is_causal, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_FNA_FORWARD_3D_SM50_float32(is_causal, cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_FNA_FORWARD_3D_SM50_float16(is_causal, cb); \ } \ else { \ @@ -137,10 +137,10 @@ namespace fna { #define DISPATCH_FNA_FORWARD_3D_SM70(dtype, is_causal, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_FNA_FORWARD_3D_SM70_float32(is_causal, cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_FNA_FORWARD_3D_SM70_float16(is_causal, cb); \ } \ else { \ @@ -150,10 +150,10 @@ namespace fna { #define DISPATCH_FNA_FORWARD_3D_SM75(dtype, is_causal, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_FNA_FORWARD_3D_SM75_float32(is_causal, cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_FNA_FORWARD_3D_SM75_float16(is_causal, cb); \ } \ else { \ @@ -163,13 +163,13 @@ namespace fna { #define DISPATCH_FNA_FORWARD_3D_SM80(dtype, is_causal, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_FNA_FORWARD_3D_SM80_float32(is_causal, cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_FNA_FORWARD_3D_SM80_float16(is_causal, cb); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ DISPATCH_FNA_FORWARD_3D_SM80_bfloat16(is_causal, cb); \ } \ else { \ @@ -179,10 +179,10 @@ namespace fna { #define DISPATCH_FNA_BACKWARD_1D_SM50(dtype, is_causal, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_FNA_BACKWARD_1D_SM50_float32(is_causal, cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_FNA_BACKWARD_1D_SM50_float16(is_causal, cb); \ } \ else { \ @@ -192,10 +192,10 @@ namespace fna { #define DISPATCH_FNA_BACKWARD_1D_SM70(dtype, is_causal, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_FNA_BACKWARD_1D_SM70_float32(is_causal, cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_FNA_BACKWARD_1D_SM70_float16(is_causal, cb); \ } \ else { \ @@ -205,10 +205,10 @@ namespace fna { #define DISPATCH_FNA_BACKWARD_1D_SM75(dtype, is_causal, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_FNA_BACKWARD_1D_SM75_float32(is_causal, cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_FNA_BACKWARD_1D_SM75_float16(is_causal, cb); \ } \ else { \ @@ -218,13 +218,13 @@ namespace fna { #define DISPATCH_FNA_BACKWARD_1D_SM80(dtype, is_causal, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_FNA_BACKWARD_1D_SM80_float32(is_causal, cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_FNA_BACKWARD_1D_SM80_float16(is_causal, cb); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ DISPATCH_FNA_BACKWARD_1D_SM80_bfloat16(is_causal, cb); \ } \ else { \ @@ -234,10 +234,10 @@ namespace fna { #define DISPATCH_FNA_BACKWARD_2D_SM50(dtype, is_causal, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_FNA_BACKWARD_2D_SM50_float32(is_causal, cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_FNA_BACKWARD_2D_SM50_float16(is_causal, cb); \ } \ else { \ @@ -247,10 +247,10 @@ namespace fna { #define DISPATCH_FNA_BACKWARD_2D_SM70(dtype, is_causal, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_FNA_BACKWARD_2D_SM70_float32(is_causal, cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_FNA_BACKWARD_2D_SM70_float16(is_causal, cb); \ } \ else { \ @@ -260,10 +260,10 @@ namespace fna { #define DISPATCH_FNA_BACKWARD_2D_SM75(dtype, is_causal, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_FNA_BACKWARD_2D_SM75_float32(is_causal, cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_FNA_BACKWARD_2D_SM75_float16(is_causal, cb); \ } \ else { \ @@ -273,13 +273,13 @@ namespace fna { #define DISPATCH_FNA_BACKWARD_2D_SM80(dtype, is_causal, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_FNA_BACKWARD_2D_SM80_float32(is_causal, cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_FNA_BACKWARD_2D_SM80_float16(is_causal, cb); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ DISPATCH_FNA_BACKWARD_2D_SM80_bfloat16(is_causal, cb); \ } \ else { \ @@ -289,10 +289,10 @@ namespace fna { #define DISPATCH_FNA_BACKWARD_3D_SM50(dtype, is_causal, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_FNA_BACKWARD_3D_SM50_float32(is_causal, cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_FNA_BACKWARD_3D_SM50_float16(is_causal, cb); \ } \ else { \ @@ -302,10 +302,10 @@ namespace fna { #define DISPATCH_FNA_BACKWARD_3D_SM70(dtype, is_causal, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_FNA_BACKWARD_3D_SM70_float32(is_causal, cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_FNA_BACKWARD_3D_SM70_float16(is_causal, cb); \ } \ else { \ @@ -315,10 +315,10 @@ namespace fna { #define DISPATCH_FNA_BACKWARD_3D_SM75(dtype, is_causal, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_FNA_BACKWARD_3D_SM75_float32(is_causal, cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_FNA_BACKWARD_3D_SM75_float16(is_causal, cb); \ } \ else { \ @@ -328,13 +328,13 @@ namespace fna { #define DISPATCH_FNA_BACKWARD_3D_SM80(dtype, is_causal, cb) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_FNA_BACKWARD_3D_SM80_float32(is_causal, cb); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_FNA_BACKWARD_3D_SM80_float16(is_causal, cb); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ DISPATCH_FNA_BACKWARD_3D_SM80_bfloat16(is_causal, cb); \ } \ else { \ diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha/interface.h b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha/interface.h index 426e3e33..c092f070 100644 --- a/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha/interface.h +++ b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha/interface.h @@ -22,10 +22,10 @@ namespace cuda { namespace fmha_hopper { #define DISPATCH_HOPPER_FMHA_FORWARD(dtype, dim, q_tile_size, kv_tile_size, kernel_type, ...) \ [&] { \ - if (dtype == torch::kFloat16) { \ + if (dtype == at::kHalf) { \ DISPATCH_HOPPER_FMHA_FORWARD_float16(dim, q_tile_size, kv_tile_size, kernel_type, __VA_ARGS__); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ DISPATCH_HOPPER_FMHA_FORWARD_bfloat16(dim, q_tile_size, kv_tile_size, kernel_type, __VA_ARGS__); \ } \ else { \ diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/interface.h b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/interface.h index a168a87f..2310231f 100644 --- a/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/interface.h +++ b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/interface.h @@ -21,10 +21,10 @@ namespace cuda { namespace fmha_hopper { #define DISPATCH_HOPPER_FMHA_BACKWARD(dtype, dim, q_tile_size, kv_tile_size, ...) \ [&] { \ - if (dtype == torch::kFloat16) { \ + if (dtype == at::kHalf) { \ DISPATCH_HOPPER_FMHA_BACKWARD_float16(dim, q_tile_size, kv_tile_size, __VA_ARGS__); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ DISPATCH_HOPPER_FMHA_BACKWARD_bfloat16(dim, q_tile_size, kv_tile_size, __VA_ARGS__); \ } \ else { \ diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_dtype.h b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_dtype.h index 385551f6..690fb01a 100644 --- a/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_dtype.h +++ b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_dtype.h @@ -22,10 +22,10 @@ namespace cuda { namespace fna_hopper { #define DISPATCH_HOPPER_FNA_FORWARD_1D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ [&] { \ - if (dtype == torch::kFloat16) { \ + if (dtype == at::kHalf) { \ DISPATCH_HOPPER_FNA_FORWARD_1D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ DISPATCH_HOPPER_FNA_FORWARD_1D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ } \ else { \ @@ -35,10 +35,10 @@ namespace fna_hopper { #define DISPATCH_HOPPER_FNA_FORWARD_2D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ [&] { \ - if (dtype == torch::kFloat16) { \ + if (dtype == at::kHalf) { \ DISPATCH_HOPPER_FNA_FORWARD_2D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ DISPATCH_HOPPER_FNA_FORWARD_2D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ } \ else { \ @@ -48,10 +48,10 @@ namespace fna_hopper { #define DISPATCH_HOPPER_FNA_FORWARD_3D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, ...) \ [&] { \ - if (dtype == torch::kFloat16) { \ + if (dtype == at::kHalf) { \ DISPATCH_HOPPER_FNA_FORWARD_3D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ DISPATCH_HOPPER_FNA_FORWARD_3D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, kernel_type, __VA_ARGS__); \ } \ else { \ diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_dtype.h b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_dtype.h index 04ff1e89..fd64a26f 100644 --- a/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_dtype.h +++ b/natten/csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_dtype.h @@ -21,10 +21,10 @@ namespace cuda { namespace fna_hopper { #define DISPATCH_HOPPER_FNA_BACKWARD_1D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, ...) \ [&] { \ - if (dtype == torch::kFloat16) { \ + if (dtype == at::kHalf) { \ DISPATCH_HOPPER_FNA_BACKWARD_1D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ DISPATCH_HOPPER_FNA_BACKWARD_1D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ } \ else { \ @@ -34,10 +34,10 @@ namespace fna_hopper { #define DISPATCH_HOPPER_FNA_BACKWARD_2D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, ...) \ [&] { \ - if (dtype == torch::kFloat16) { \ + if (dtype == at::kHalf) { \ DISPATCH_HOPPER_FNA_BACKWARD_2D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ DISPATCH_HOPPER_FNA_BACKWARD_2D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ } \ else { \ @@ -47,10 +47,10 @@ namespace fna_hopper { #define DISPATCH_HOPPER_FNA_BACKWARD_3D(dtype, dim, is_causal, q_tile_shape, kv_tile_shape, ...) \ [&] { \ - if (dtype == torch::kFloat16) { \ + if (dtype == at::kHalf) { \ DISPATCH_HOPPER_FNA_BACKWARD_3D_float16(dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ DISPATCH_HOPPER_FNA_BACKWARD_3D_bfloat16(dim, is_causal, q_tile_shape, kv_tile_shape, __VA_ARGS__); \ } \ else { \ diff --git a/natten/csrc/autogen/include/natten_autogen/cuda/reference/dispatch_dtype.h b/natten/csrc/autogen/include/natten_autogen/cuda/reference/dispatch_dtype.h index fe05f1b4..ce9e621b 100644 --- a/natten/csrc/autogen/include/natten_autogen/cuda/reference/dispatch_dtype.h +++ b/natten/csrc/autogen/include/natten_autogen/cuda/reference/dispatch_dtype.h @@ -20,13 +20,13 @@ namespace cuda { namespace reference { #define DISPATCH_REFERENCE_FNA_FORWARD_1D(dtype, is_causal, ...) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_REFERENCE_FNA_FORWARD_1D_float32(is_causal, __VA_ARGS__); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_REFERENCE_FNA_FORWARD_1D_float16(is_causal, __VA_ARGS__); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ DISPATCH_REFERENCE_FNA_FORWARD_1D_bfloat16(is_causal, __VA_ARGS__); \ } \ else { \ @@ -36,13 +36,13 @@ namespace reference { #define DISPATCH_REFERENCE_FNA_BACKWARD_1D(dtype, is_causal, ...) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_REFERENCE_FNA_BACKWARD_1D_float32(is_causal, __VA_ARGS__); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_REFERENCE_FNA_BACKWARD_1D_float16(is_causal, __VA_ARGS__); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ DISPATCH_REFERENCE_FNA_BACKWARD_1D_bfloat16(is_causal, __VA_ARGS__); \ } \ else { \ @@ -52,13 +52,13 @@ namespace reference { #define DISPATCH_REFERENCE_FNA_FORWARD_2D(dtype, is_causal, ...) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_REFERENCE_FNA_FORWARD_2D_float32(is_causal, __VA_ARGS__); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_REFERENCE_FNA_FORWARD_2D_float16(is_causal, __VA_ARGS__); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ DISPATCH_REFERENCE_FNA_FORWARD_2D_bfloat16(is_causal, __VA_ARGS__); \ } \ else { \ @@ -68,13 +68,13 @@ namespace reference { #define DISPATCH_REFERENCE_FNA_BACKWARD_2D(dtype, is_causal, ...) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_REFERENCE_FNA_BACKWARD_2D_float32(is_causal, __VA_ARGS__); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_REFERENCE_FNA_BACKWARD_2D_float16(is_causal, __VA_ARGS__); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ DISPATCH_REFERENCE_FNA_BACKWARD_2D_bfloat16(is_causal, __VA_ARGS__); \ } \ else { \ @@ -84,13 +84,13 @@ namespace reference { #define DISPATCH_REFERENCE_FNA_FORWARD_3D(dtype, is_causal, ...) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_REFERENCE_FNA_FORWARD_3D_float32(is_causal, __VA_ARGS__); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_REFERENCE_FNA_FORWARD_3D_float16(is_causal, __VA_ARGS__); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ DISPATCH_REFERENCE_FNA_FORWARD_3D_bfloat16(is_causal, __VA_ARGS__); \ } \ else { \ @@ -100,13 +100,13 @@ namespace reference { #define DISPATCH_REFERENCE_FNA_BACKWARD_3D(dtype, is_causal, ...) \ [&] { \ - if (dtype == torch::kFloat32) { \ + if (dtype == at::kFloat) { \ DISPATCH_REFERENCE_FNA_BACKWARD_3D_float32(is_causal, __VA_ARGS__); \ } \ - else if (dtype == torch::kFloat16) { \ + else if (dtype == at::kHalf) { \ DISPATCH_REFERENCE_FNA_BACKWARD_3D_float16(is_causal, __VA_ARGS__); \ } \ - else if (dtype == torch::kBFloat16) { \ + else if (dtype == at::kBFloat16) { \ DISPATCH_REFERENCE_FNA_BACKWARD_3D_bfloat16(is_causal, __VA_ARGS__); \ } \ else { \ diff --git a/natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_0.cu b/natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_0.cu index 1e08b3cc..4e0be56b 100644 --- a/natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_0.cu +++ b/natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_0.cu @@ -75,7 +75,7 @@ void blackwell_fmha_backward_float16_128x128x32( if constexpr (IsDeterministic) { auto kBlockM = cute::get<0>(GemmShape{}); int effective_seqlen_q = is_varlen ? max_seqlen_Q : seqlen_q; - dq_semaphore = torch::zeros( + dq_semaphore = at::zeros( {(effective_seqlen_q + kBlockM - 1) / kBlockM, batch_size, heads_q}, tensor_options.dtype(at::ScalarType::Int)); dq_semaphore_ptr = static_cast(dq_semaphore.data_ptr()); @@ -208,7 +208,7 @@ void blackwell_fmha_backward_float16_128x128x64( if constexpr (IsDeterministic) { auto kBlockM = cute::get<0>(GemmShape{}); int effective_seqlen_q = is_varlen ? max_seqlen_Q : seqlen_q; - dq_semaphore = torch::zeros( + dq_semaphore = at::zeros( {(effective_seqlen_q + kBlockM - 1) / kBlockM, batch_size, heads_q}, tensor_options.dtype(at::ScalarType::Int)); dq_semaphore_ptr = static_cast(dq_semaphore.data_ptr()); diff --git a/natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_1.cu b/natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_1.cu index 9cb855d9..e10b16b9 100644 --- a/natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_1.cu +++ b/natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_1.cu @@ -75,7 +75,7 @@ void blackwell_fmha_backward_float16_128x128x128( if constexpr (IsDeterministic) { auto kBlockM = cute::get<0>(GemmShape{}); int effective_seqlen_q = is_varlen ? max_seqlen_Q : seqlen_q; - dq_semaphore = torch::zeros( + dq_semaphore = at::zeros( {(effective_seqlen_q + kBlockM - 1) / kBlockM, batch_size, heads_q}, tensor_options.dtype(at::ScalarType::Int)); dq_semaphore_ptr = static_cast(dq_semaphore.data_ptr()); @@ -208,7 +208,7 @@ void blackwell_fmha_backward_bfloat16_128x128x32( if constexpr (IsDeterministic) { auto kBlockM = cute::get<0>(GemmShape{}); int effective_seqlen_q = is_varlen ? max_seqlen_Q : seqlen_q; - dq_semaphore = torch::zeros( + dq_semaphore = at::zeros( {(effective_seqlen_q + kBlockM - 1) / kBlockM, batch_size, heads_q}, tensor_options.dtype(at::ScalarType::Int)); dq_semaphore_ptr = static_cast(dq_semaphore.data_ptr()); diff --git a/natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_2.cu b/natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_2.cu index 0ebf4405..60599df2 100644 --- a/natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_2.cu +++ b/natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_2.cu @@ -75,7 +75,7 @@ void blackwell_fmha_backward_bfloat16_128x128x64( if constexpr (IsDeterministic) { auto kBlockM = cute::get<0>(GemmShape{}); int effective_seqlen_q = is_varlen ? max_seqlen_Q : seqlen_q; - dq_semaphore = torch::zeros( + dq_semaphore = at::zeros( {(effective_seqlen_q + kBlockM - 1) / kBlockM, batch_size, heads_q}, tensor_options.dtype(at::ScalarType::Int)); dq_semaphore_ptr = static_cast(dq_semaphore.data_ptr()); diff --git a/natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_3.cu b/natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_3.cu index b934754f..621d25ff 100644 --- a/natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_3.cu +++ b/natten/csrc/autogen/src/cuda/blackwell_fmha_bwd/source_3.cu @@ -75,7 +75,7 @@ void blackwell_fmha_backward_bfloat16_128x128x128( if constexpr (IsDeterministic) { auto kBlockM = cute::get<0>(GemmShape{}); int effective_seqlen_q = is_varlen ? max_seqlen_Q : seqlen_q; - dq_semaphore = torch::zeros( + dq_semaphore = at::zeros( {(effective_seqlen_q + kBlockM - 1) / kBlockM, batch_size, heads_q}, tensor_options.dtype(at::ScalarType::Int)); dq_semaphore_ptr = static_cast(dq_semaphore.data_ptr()); diff --git a/natten/csrc/include/natten/helpers.h b/natten/csrc/include/natten/helpers.h index 05521770..3c7c1055 100644 --- a/natten/csrc/include/natten/helpers.h +++ b/natten/csrc/include/natten/helpers.h @@ -316,7 +316,7 @@ void CheckLogSumExp(const at::Tensor& output, const at::Tensor& logsumexp) { // Logsumexp: [batch, *, heads] static_assert(NaDim >= 1 && NaDim < 4); TORCH_CHECK( - logsumexp.scalar_type() == torch::kFloat, + logsumexp.scalar_type() == at::kFloat, "`logsumexp` must be stored in float32 data type, got ", logsumexp.scalar_type()); TORCH_CHECK( @@ -357,7 +357,7 @@ inline void CheckLogSumExpHeadsFirst( // Output: [batch, seqlen, heads, dim] // Logsumexp: [batch, heads, seqlen] TORCH_CHECK( - logsumexp.scalar_type() == torch::kFloat, + logsumexp.scalar_type() == at::kFloat, "`logsumexp` must be stored in float32 data type, got ", logsumexp.scalar_type()); TORCH_CHECK( @@ -416,9 +416,9 @@ inline void AssertDimsAre128BitAligned( "."); TORCH_CHECK( - query.scalar_type() == torch::kFloat || - query.scalar_type() == torch::kFloat16 || - query.scalar_type() == torch::kBFloat16 || + query.scalar_type() == at::kFloat || + query.scalar_type() == at::kHalf || + query.scalar_type() == at::kBFloat16 || query.scalar_type() == c10::ScalarType::Float8_e4m3fn || query.scalar_type() == c10::ScalarType::Float8_e5m2, "This NATTEN operation only supports FP32, FP16, BF16, and FP8 data types, got ", @@ -426,11 +426,11 @@ inline void AssertDimsAre128BitAligned( "."); int alignment; - if (query.scalar_type() == torch::kFloat) { + if (query.scalar_type() == at::kFloat) { alignment = 4; } else if ( - query.scalar_type() == torch::kFloat16 || - query.scalar_type() == torch::kBFloat16) { + query.scalar_type() == at::kHalf || + query.scalar_type() == at::kBFloat16) { alignment = 8; } else if ( query.scalar_type() == c10::ScalarType::Float8_e4m3fn || diff --git a/natten/scripts/autogen_blackwell_fmha.py b/natten/scripts/autogen_blackwell_fmha.py index 8dd359b1..60604735 100644 --- a/natten/scripts/autogen_blackwell_fmha.py +++ b/natten/scripts/autogen_blackwell_fmha.py @@ -160,8 +160,8 @@ def __init__(self, name, short_name, torch_name, bits): self.torch_name = torch_name -Half = DataType("cutlass::half_t", "float16", "torch::kFloat16", 16) -BFloat = DataType("cutlass::bfloat16_t", "bfloat16", "torch::kBFloat16", 16) +Half = DataType("cutlass::half_t", "float16", "at::kHalf", 16) +BFloat = DataType("cutlass::bfloat16_t", "bfloat16", "at::kBFloat16", 16) E4M3 = DataType("cutlass::float_e4m3_t", "e4m3", "c10::ScalarType::Float8_e4m3fn", 8) E5M2 = DataType("cutlass::float_e5m2_t", "e5m2", "c10::ScalarType::Float8_e5m2", 8) diff --git a/natten/scripts/autogen_blackwell_fmha_bwd.py b/natten/scripts/autogen_blackwell_fmha_bwd.py index 9a6ba0d3..ed57eac7 100644 --- a/natten/scripts/autogen_blackwell_fmha_bwd.py +++ b/natten/scripts/autogen_blackwell_fmha_bwd.py @@ -114,7 +114,7 @@ if constexpr (IsDeterministic) {{ auto kBlockM = cute::get<0>(GemmShape{{}}); int effective_seqlen_q = is_varlen ? max_seqlen_Q : seqlen_q; - dq_semaphore = torch::zeros( + dq_semaphore = at::zeros( {{(effective_seqlen_q + kBlockM - 1) / kBlockM, batch_size, heads_q}}, tensor_options.dtype(at::ScalarType::Int)); dq_semaphore_ptr = static_cast(dq_semaphore.data_ptr()); @@ -195,8 +195,8 @@ def __init__(self, name, short_name, torch_name, bits): self.torch_name = torch_name -Half = DataType("cutlass::half_t", "float16", "torch::kFloat16", 16) -BFloat = DataType("cutlass::bfloat16_t", "bfloat16", "torch::kBFloat16", 16) +Half = DataType("cutlass::half_t", "float16", "at::kHalf", 16) +BFloat = DataType("cutlass::bfloat16_t", "bfloat16", "at::kBFloat16", 16) def iterable_to_static_cute_tuple(shape_in) -> str: diff --git a/natten/scripts/autogen_blackwell_fna.py b/natten/scripts/autogen_blackwell_fna.py index 502fd365..32df0b67 100644 --- a/natten/scripts/autogen_blackwell_fna.py +++ b/natten/scripts/autogen_blackwell_fna.py @@ -123,8 +123,8 @@ def __init__(self, name, short_name, torch_name, bits): self.torch_name = torch_name -Half = DataType("cutlass::half_t", "float16", "torch::kFloat16", 16) -BFloat = DataType("cutlass::bfloat16_t", "bfloat16", "torch::kBFloat16", 16) +Half = DataType("cutlass::half_t", "float16", "at::kHalf", 16) +BFloat = DataType("cutlass::bfloat16_t", "bfloat16", "at::kBFloat16", 16) E4M3 = DataType("cutlass::float_e4m3_t", "e4m3", "c10::ScalarType::Float8_e4m3fn", 8) E5M2 = DataType("cutlass::float_e5m2_t", "e5m2", "c10::ScalarType::Float8_e5m2", 8) diff --git a/natten/scripts/autogen_blackwell_fna_bwd.py b/natten/scripts/autogen_blackwell_fna_bwd.py index 08cbfbbd..c7379df3 100644 --- a/natten/scripts/autogen_blackwell_fna_bwd.py +++ b/natten/scripts/autogen_blackwell_fna_bwd.py @@ -134,8 +134,8 @@ def __init__(self, name, short_name, torch_name, bits): self.torch_name = torch_name -Half = DataType("cutlass::half_t", "float16", "torch::kFloat16", 16) -BFloat = DataType("cutlass::bfloat16_t", "bfloat16", "torch::kBFloat16", 16) +Half = DataType("cutlass::half_t", "float16", "at::kHalf", 16) +BFloat = DataType("cutlass::bfloat16_t", "bfloat16", "at::kBFloat16", 16) def iterable_to_static_cute_tuple(shape_in) -> str: diff --git a/natten/scripts/autogen_fmha.py b/natten/scripts/autogen_fmha.py index d30b4062..bcba7f14 100644 --- a/natten/scripts/autogen_fmha.py +++ b/natten/scripts/autogen_fmha.py @@ -69,9 +69,9 @@ def __init__(self, name, torch_name, short_name, bits, min_sm): self.min_sm = min_sm -NATTEN_Float = DataType("float", "torch::kFloat32", "float32", 32, 0) -NATTEN_Half = DataType("cutlass::half_t", "torch::kFloat16", "float16", 16, 50) -NATTEN_BFloat = DataType("cutlass::bfloat16_t", "torch::kBFloat16", "bfloat16", 16, 80) +NATTEN_Float = DataType("float", "at::kFloat", "float32", 32, 0) +NATTEN_Half = DataType("cutlass::half_t", "at::kHalf", "float16", 16, 50) +NATTEN_BFloat = DataType("cutlass::bfloat16_t", "at::kBFloat16", "bfloat16", 16, 80) KERNEL_DECL_TEMPLATE = """__global__ void __launch_bounds__( diff --git a/natten/scripts/autogen_fna.py b/natten/scripts/autogen_fna.py index 09fdcc05..e709228e 100644 --- a/natten/scripts/autogen_fna.py +++ b/natten/scripts/autogen_fna.py @@ -73,9 +73,9 @@ def __init__(self, name, torch_name, short_name, bits, min_sm): self.min_sm = min_sm -NATTEN_Float = DataType("float", "torch::kFloat32", "float32", 32, 0) -NATTEN_Half = DataType("cutlass::half_t", "torch::kFloat16", "float16", 16, 50) -NATTEN_BFloat = DataType("cutlass::bfloat16_t", "torch::kBFloat16", "bfloat16", 16, 80) +NATTEN_Float = DataType("float", "at::kFloat", "float32", 32, 0) +NATTEN_Half = DataType("cutlass::half_t", "at::kHalf", "float16", 16, 50) +NATTEN_BFloat = DataType("cutlass::bfloat16_t", "at::kBFloat16", "bfloat16", 16, 80) KERNEL_DECL_TEMPLATE = """__global__ void __launch_bounds__( diff --git a/natten/scripts/autogen_hopper_fmha.py b/natten/scripts/autogen_hopper_fmha.py index 29769c81..3c3ccb02 100644 --- a/natten/scripts/autogen_hopper_fmha.py +++ b/natten/scripts/autogen_hopper_fmha.py @@ -186,8 +186,8 @@ def __init__(self, name, short_name, torch_name, bits): self.torch_name = torch_name -Half = DataType("cutlass::half_t", "float16", "torch::kFloat16", 16) -BFloat = DataType("cutlass::bfloat16_t", "bfloat16", "torch::kBFloat16", 16) +Half = DataType("cutlass::half_t", "float16", "at::kHalf", 16) +BFloat = DataType("cutlass::bfloat16_t", "bfloat16", "at::kBFloat16", 16) def iterable_to_static_cute_tuple(shape_in) -> str: diff --git a/natten/scripts/autogen_hopper_fmha_bwd.py b/natten/scripts/autogen_hopper_fmha_bwd.py index 9ecbbb48..c19c8941 100644 --- a/natten/scripts/autogen_hopper_fmha_bwd.py +++ b/natten/scripts/autogen_hopper_fmha_bwd.py @@ -177,8 +177,8 @@ def __init__(self, name, short_name, torch_name, bits): self.torch_name = torch_name -Half = DataType("cutlass::half_t", "float16", "torch::kFloat16", 16) -BFloat = DataType("cutlass::bfloat16_t", "bfloat16", "torch::kBFloat16", 16) +Half = DataType("cutlass::half_t", "float16", "at::kHalf", 16) +BFloat = DataType("cutlass::bfloat16_t", "bfloat16", "at::kBFloat16", 16) def iterable_to_static_cute_tuple(shape_in) -> str: diff --git a/natten/scripts/autogen_hopper_fna.py b/natten/scripts/autogen_hopper_fna.py index e7504a43..84c44d34 100644 --- a/natten/scripts/autogen_hopper_fna.py +++ b/natten/scripts/autogen_hopper_fna.py @@ -145,8 +145,8 @@ def __init__(self, name, short_name, torch_name, bits): self.torch_name = torch_name -Half = DataType("cutlass::half_t", "float16", "torch::kFloat16", 16) -BFloat = DataType("cutlass::bfloat16_t", "bfloat16", "torch::kBFloat16", 16) +Half = DataType("cutlass::half_t", "float16", "at::kHalf", 16) +BFloat = DataType("cutlass::bfloat16_t", "bfloat16", "at::kBFloat16", 16) def iterable_to_static_cute_tuple(shape_in) -> str: diff --git a/natten/scripts/autogen_hopper_fna_bwd.py b/natten/scripts/autogen_hopper_fna_bwd.py index d3d4770d..7ec37719 100644 --- a/natten/scripts/autogen_hopper_fna_bwd.py +++ b/natten/scripts/autogen_hopper_fna_bwd.py @@ -137,8 +137,8 @@ def __init__(self, name, short_name, torch_name, bits): self.torch_name = torch_name -Half = DataType("cutlass::half_t", "float16", "torch::kFloat16", 16) -BFloat = DataType("cutlass::bfloat16_t", "bfloat16", "torch::kBFloat16", 16) +Half = DataType("cutlass::half_t", "float16", "at::kHalf", 16) +BFloat = DataType("cutlass::bfloat16_t", "bfloat16", "at::kBFloat16", 16) def iterable_to_static_cute_tuple(shape_in) -> str: diff --git a/natten/scripts/autogen_reference_fna.py b/natten/scripts/autogen_reference_fna.py index 077bfd73..edc200e2 100644 --- a/natten/scripts/autogen_reference_fna.py +++ b/natten/scripts/autogen_reference_fna.py @@ -176,9 +176,9 @@ def __init__(self, name, short_name, torch_name, bits): self.torch_name = torch_name -Float = DataType("float", "float32", "torch::kFloat32", 32) -Half = DataType("cutlass::half_t", "float16", "torch::kFloat16", 16) -BFloat = DataType("cutlass::bfloat16_t", "bfloat16", "torch::kBFloat16", 16) +Float = DataType("float", "float32", "at::kFloat", 32) +Half = DataType("cutlass::half_t", "float16", "at::kHalf", 16) +BFloat = DataType("cutlass::bfloat16_t", "bfloat16", "at::kBFloat16", 16) def iterable_to_static_cute_tuple(shape_in) -> str: From 5a16a003287605036597e962c359f7545e593c7b Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Fri, 17 Jul 2026 08:08:18 +0530 Subject: [PATCH 5/8] natten: rename types.py to _types.py to avoid shadowing stdlib kernel-builder installs the package contents flat into the build variant directory, and test runners put that directory on PYTHONPATH. natten's top-level types module therefore shadowed the standard library's types module, and Python failed during site initialization before any test ran (https://github.com/huggingface/kernels-community/actions/runs/29493798395). Rename the module to _types (the only stdlib-colliding name in the package), rewrite the 28 internal import sites, and alias it back as natten.types in __init__ so the upstream import surface (from natten.types import ...) is unchanged, including for the vendored tests. Co-Authored-By: Claude Fable 5 --- natten/CMakeLists.txt | 662 ++++++++++++++++++ natten/cmake/_ops.py.in | 9 + natten/cmake/add_build_metadata.py | 56 ++ natten/cmake/build-variants.cmake | 328 +++++++++ natten/cmake/compile-metal.cmake | 104 +++ natten/cmake/get_gpu_lang.cmake | 9 + natten/cmake/get_gpu_lang.py | 20 + natten/cmake/hipify.py | 76 ++ natten/cmake/kernel.cmake | 317 +++++++++ natten/cmake/metallib_to_header.py | 73 ++ natten/cmake/utils.cmake | 611 ++++++++++++++++ natten/compat.py | 26 + natten/metadata-cann.json | 22 + natten/metadata-cpu.json | 22 + natten/metadata-cuda.json | 22 + natten/metadata-metal.json | 22 + natten/metadata-neuron.json | 22 + natten/metadata-rocm.json | 22 + natten/metadata-xpu.json | 22 + natten/pyproject.toml | 23 + natten/setup.py | 200 ++++++ natten/torch-ext/natten/__init__.py | 12 + .../torch-ext/natten/{types.py => _types.py} | 0 .../natten/backends/blackwell_fmha.py | 2 +- .../natten/backends/blackwell_fna.py | 2 +- .../natten/backends/configs/__init__.py | 2 +- .../backends/configs/cutlass/__init__.py | 2 +- .../configs/cutlass/backward_knobs.py | 2 +- .../configs/cutlass/fna_backward_128x128.py | 2 +- .../configs/cutlass/fna_backward_128x64.py | 2 +- .../configs/cutlass/fna_backward_64x64.py | 2 +- .../configs/cutlass/fna_forward_32x128.py | 2 +- .../configs/cutlass/fna_forward_64x128.py | 2 +- .../configs/cutlass/fna_forward_64x64.py | 2 +- .../configs/cutlass_blackwell/__init__.py | 2 +- .../configs/cutlass_hopper/__init__.py | 2 +- .../natten/backends/configs/flex/__init__.py | 2 +- natten/torch-ext/natten/backends/flex.py | 2 +- natten/torch-ext/natten/backends/fmha.py | 2 +- natten/torch-ext/natten/backends/fna.py | 2 +- .../torch-ext/natten/backends/hopper_fmha.py | 2 +- .../torch-ext/natten/backends/hopper_fna.py | 2 +- natten/torch-ext/natten/backends/reference.py | 2 +- natten/torch-ext/natten/functional.py | 2 +- natten/torch-ext/natten/modules.py | 2 +- .../natten/token_permute/cutlass_impl.py | 2 +- .../natten/token_permute/frontend.py | 2 +- .../natten/token_permute/torch_impl.py | 2 +- natten/torch-ext/natten/utils/checks.py | 2 +- natten/torch-ext/natten/utils/tuples.py | 2 +- natten/torch-ext/natten/utils/varlen.py | 2 +- natten/torch-ext/registration.h | 39 ++ 52 files changed, 2747 insertions(+), 28 deletions(-) create mode 100644 natten/CMakeLists.txt create mode 100644 natten/cmake/_ops.py.in create mode 100644 natten/cmake/add_build_metadata.py create mode 100644 natten/cmake/build-variants.cmake create mode 100644 natten/cmake/compile-metal.cmake create mode 100644 natten/cmake/get_gpu_lang.cmake create mode 100644 natten/cmake/get_gpu_lang.py create mode 100644 natten/cmake/hipify.py create mode 100644 natten/cmake/kernel.cmake create mode 100644 natten/cmake/metallib_to_header.py create mode 100644 natten/cmake/utils.cmake create mode 100644 natten/compat.py create mode 100644 natten/metadata-cann.json create mode 100644 natten/metadata-cpu.json create mode 100644 natten/metadata-cuda.json create mode 100644 natten/metadata-metal.json create mode 100644 natten/metadata-neuron.json create mode 100644 natten/metadata-rocm.json create mode 100644 natten/metadata-xpu.json create mode 100644 natten/pyproject.toml create mode 100644 natten/setup.py rename natten/torch-ext/natten/{types.py => _types.py} (100%) create mode 100644 natten/torch-ext/registration.h diff --git a/natten/CMakeLists.txt b/natten/CMakeLists.txt new file mode 100644 index 00000000..cbf24aec --- /dev/null +++ b/natten/CMakeLists.txt @@ -0,0 +1,662 @@ +cmake_minimum_required(VERSION 3.26) + +# Set Intel SYCL compiler before project() call +find_program(ICX_COMPILER icx) +find_program(ICPX_COMPILER icpx) + +if(ICX_COMPILER OR ICPX_COMPILER) + set(CMAKE_C_COMPILER ${ICX_COMPILER}) + + if(WIN32) + set(CMAKE_CXX_COMPILER ${ICX_COMPILER}) + else() + set(CMAKE_CXX_COMPILER ${ICPX_COMPILER}) + endif() +endif() + +project( LANGUAGES CXX) + +install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS) + +include(FetchContent) +file(MAKE_DIRECTORY ${FETCHCONTENT_BASE_DIR}) # Ensure the directory exists +message(STATUS "FetchContent base directory: ${FETCHCONTENT_BASE_DIR}") + +set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1200;gfx1201") + +# Make Torch CMake machinery happy. Whatever we use does not matter, since +# we set the arches per-file later anyway. +set(ENV{PYTORCH_ROCM_ARCH} "${HIP_SUPPORTED_ARCHS}") + +include(${CMAKE_CURRENT_LIST_DIR}/cmake/utils.cmake) +include(${CMAKE_CURRENT_LIST_DIR}/cmake/kernel.cmake) +include(${CMAKE_CURRENT_LIST_DIR}/cmake/get_gpu_lang.cmake) + +if(DEFINED Python3_EXECUTABLE) + # Allow passing through the interpreter (e.g. from setup.py). + find_package(Python3 COMPONENTS Development Development.SABIModule Interpreter) + if (NOT Python3_FOUND) + message(FATAL_ERROR "Unable to find python matching: ${EXECUTABLE}.") + endif() +else() + find_package(Python3 REQUIRED COMPONENTS Development Development.SABIModule Interpreter) +endif() + +get_gpu_lang(DETECTED_GPU_LANG) +set(GPU_LANG "${DETECTED_GPU_LANG}" CACHE STRING "GPU language") +gpu_lang_to_backend(BACKEND "${GPU_LANG}") +message(STATUS "Using backend: ${BACKEND}, GPU language: ${GPU_LANG}") + +set(OPS_NAME "_natten_${BACKEND}_28fa1dd") + +append_cmake_prefix_path("torch" "torch.utils.cmake_prefix_path") + +find_package(Torch REQUIRED) + +run_python(TORCH_VERSION "import torch; print(torch.__version__.split('+')[0])" "Failed to get Torch version") + + + + +option(BUILD_ALL_SUPPORTED_ARCHS "Build all supported architectures" off) + +if(DEFINED CMAKE_CUDA_COMPILER_VERSION AND + CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0) + set(CUDA_DEFAULT_KERNEL_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;11.0;12.0;12.1+PTX") +elseif(DEFINED CMAKE_CUDA_COMPILER_VERSION AND + CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.8) + set(CUDA_DEFAULT_KERNEL_ARCHS "7.0;7.2;7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.1;12.0+PTX") +else() + set(CUDA_DEFAULT_KERNEL_ARCHS "7.0;7.2;7.5;8.0;8.6;8.7;8.9;9.0+PTX") +endif() + +# Basic checks for each GPU language. +if(GPU_LANG STREQUAL "CUDA") + if(NOT CUDA_FOUND) + message(FATAL_ERROR "GPU language is set to CUDA, but cannot find CUDA toolkit") + endif() + + + + # This clears out -gencode arguments from `CMAKE_CUDA_FLAGS`, which we need + # to set our own set of capabilities. + clear_gencode_flags() + + # Get the capabilities without +PTX suffixes, so that we can use them as + # the target archs in the loose intersection with a kernel's capabilities. + cuda_remove_ptx_suffixes(CUDA_ARCHS "${CUDA_DEFAULT_KERNEL_ARCHS}") + message(STATUS "CUDA base archs used for intersection with kernel archs: ${CUDA_ARCHS}") + + if(BUILD_ALL_SUPPORTED_ARCHS) + set(CUDA_KERNEL_ARCHS "${CUDA_DEFAULT_KERNEL_ARCHS}") + else() + try_run_python(CUDA_KERNEL_ARCHS SUCCESS "import torch; cc=torch.cuda.get_device_capability(); print(f\"{cc[0]}.{cc[1]}\")" "Failed to get CUDA capability") + if(NOT SUCCESS) + message(WARNING "Failed to detect CUDA capability, using default capabilities.") + set(CUDA_KERNEL_ARCHS "${CUDA_DEFAULT_KERNEL_ARCHS}") + endif() + endif() + + message(STATUS "Default CUDA kernel architectures: ${CUDA_KERNEL_ARCHS}") + + if(NVCC_THREADS AND GPU_LANG STREQUAL "CUDA") + message(STATUS "Using nvcc with: -threads=${NVCC_THREADS}") + list(APPEND GPU_FLAGS "--threads=${NVCC_THREADS}") + endif() + + # TODO: deprecate one of these settings. + add_compile_definitions(USE_CUDA=1) + add_compile_definitions(CUDA_KERNEL) +elseif(GPU_LANG STREQUAL "HIP") + if(NOT HIP_FOUND AND NOT PYTORCH_FOUND_HIP) + message(FATAL_ERROR "GPU language is set to HIP, but cannot find ROCm toolkit") + endif() + + # Importing torch recognizes and sets up some HIP/ROCm configuration but does + # not let cmake recognize .hip files. In order to get cmake to understand the + # .hip extension automatically, HIP must be enabled explicitly. + enable_language(HIP) + + set(ROCM_ARCHS ${HIP_SUPPORTED_ARCHS}) + message(STATUS "ROCM supported target architectures: ${ROCM_ARCHS}") + + # TODO: deprecate one of these settings. + add_compile_definitions(USE_ROCM=1) + add_compile_definitions(ROCM_KERNEL) +elseif(GPU_LANG STREQUAL "CPU") + add_compile_definitions(CPU_KERNEL) + set(CMAKE_OSX_DEPLOYMENT_TARGET "15.0" CACHE STRING "Minimum macOS deployment version") +elseif(GPU_LANG STREQUAL "METAL") + set(CMAKE_OSX_DEPLOYMENT_TARGET "26.0" CACHE STRING "Minimum macOS deployment version") + enable_language(C OBJC OBJCXX) + + add_compile_definitions(METAL_KERNEL) + + # Initialize lists for Metal shader sources and their include directories + set(ALL_METAL_SOURCES) + set(METAL_INCLUDE_DIRS) +elseif(GPU_LANG STREQUAL "SYCL") + if(NOT ICX_COMPILER AND NOT ICPX_COMPILER) + message(FATAL_ERROR "Intel SYCL C++ compiler (icpx) and/or C compiler (icx) not found. Please install Intel oneAPI toolkit.") + endif() + + execute_process( + COMMAND ${ICPX_COMPILER} --version + OUTPUT_VARIABLE ICPX_VERSION_OUTPUT + OUTPUT_STRIP_TRAILING_WHITESPACE + ) + string(REGEX MATCH "[0-9]+\\.[0-9]+" DPCPP_VERSION "${ICPX_VERSION_OUTPUT}") + set(DPCPP_VERSION "${DPCPP_VERSION}" CACHE STRING "DPCPP major.minor version") + + # On Windows, use icx (MSVC-compatible) for C++ to work with Ninja generator + # On Linux, use icpx (GNU-compatible) for C++ + if(WIN32) + message(STATUS "Using Intel SYCL C++ compiler: ${ICX_COMPILER} and C compiler: ${ICX_COMPILER} Version: ${DPCPP_VERSION} (Windows MSVC-compatible mode)") + else() + message(STATUS "Using Intel SYCL C++ compiler: ${ICPX_COMPILER} and C compiler: ${ICX_COMPILER} Version: ${DPCPP_VERSION}") + endif() + + + set(sycl_link_flags "-Wl,-z,noexecstack;-fsycl;--offload-compress;-fsycl-targets=spir64_gen,spir64;-Xs;-device pvc,xe-lpg,ats-m150 -options ' -cl-intel-enable-auto-large-GRF-mode -cl-poison-unsupported-fp64-kernels -cl-intel-greater-than-4GB-buffer-required';") + set(sycl_flags "-fPIC;-fsycl;-fhonor-nans;-fhonor-infinities;-fno-associative-math;-fno-approx-func;-fno-sycl-instrument-device-code;--offload-compress;-fsycl-targets=spir64_gen,spir64;") + set(GPU_FLAGS "${sycl_flags}") + + + add_compile_definitions(XPU_KERNEL) + add_compile_definitions(USE_XPU) +else() + message(FATAL_ERROR "Unsupported GPU language: ${GPU_LANG}") +endif() + +# Initialize SRC list for kernel and binding sources +set(SRC "") + +include(${CMAKE_CURRENT_LIST_DIR}/cmake/build-variants.cmake) + +# Generate build variant name. +if(GPU_LANG STREQUAL "CUDA") + generate_build_name(BUILD_VARIANT_NAME "${TORCH_VERSION}" "cuda" "${CUDA_VERSION}" ${_STABLE_ABI_ARG}) +elseif(GPU_LANG STREQUAL "HIP") + run_python(ROCM_VERSION "import torch.version; print(torch.version.hip.split('.')[0] + '.' + torch.version.hip.split('.')[1])" "Failed to get ROCm version") + generate_build_name(BUILD_VARIANT_NAME "${TORCH_VERSION}" "rocm" "${ROCM_VERSION}" ${_STABLE_ABI_ARG}) +elseif(GPU_LANG STREQUAL "SYCL") + generate_build_name(BUILD_VARIANT_NAME "${TORCH_VERSION}" "xpu" "${DPCPP_VERSION}" ${_STABLE_ABI_ARG}) +elseif(GPU_LANG STREQUAL "METAL") + generate_build_name(BUILD_VARIANT_NAME "${TORCH_VERSION}" "metal" "" ${_STABLE_ABI_ARG}) +elseif(GPU_LANG STREQUAL "CPU") + generate_build_name(BUILD_VARIANT_NAME "${TORCH_VERSION}" "cpu" "" ${_STABLE_ABI_ARG}) +else() + message(FATAL_ERROR "Cannot generate build name for unknown GPU_LANG: ${GPU_LANG}") +endif() + +configure_file( + ${CMAKE_CURRENT_LIST_DIR}/cmake/_ops.py.in + ${CMAKE_CURRENT_SOURCE_DIR}/torch-ext/natten/_ops.py + @ONLY +) + +if(GPU_LANG STREQUAL "CUDA") + message(STATUS "Including CUTLASS dependency") + +find_package(NvidiaCutlass) + +if (NOT NvidiaCutlass_FOUND) + set(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library") + +# Set CUTLASS_REVISION manually -- its revision detection doesn't work in this case. + set(CUTLASS_REVISION "v4.5.2" CACHE STRING "CUTLASS revision to use") + + +# Use the specified CUTLASS source directory for compilation if CUTLASS_SRC_DIR is provided + if (DEFINED ENV{CUTLASS_SRC_DIR}) + set(CUTLASS_SRC_DIR $ENV{CUTLASS_SRC_DIR}) + endif() + + if(CUTLASS_SRC_DIR) + if(NOT IS_ABSOLUTE CUTLASS_SRC_DIR) + get_filename_component(CUTLASS_SRC_DIR "${CUTLASS_SRC_DIR}" ABSOLUTE) + endif() + message(STATUS "The CUTLASS_SRC_DIR is set, using ${CUTLASS_SRC_DIR} for compilation") + FetchContent_Declare(cutlass SOURCE_DIR ${CUTLASS_SRC_DIR}) + else() + FetchContent_Declare( + cutlass + GIT_REPOSITORY https://github.com/nvidia/cutlass.git + GIT_TAG ${CUTLASS_REVISION} + GIT_PROGRESS TRUE + + # Speed up CUTLASS download by retrieving only the specified GIT_TAG instead of the history. + # Important: If GIT_SHALLOW is enabled then GIT_TAG works only with branch names and tags. + # So if the GIT_TAG above is updated to a commit hash, GIT_SHALLOW must be set to FALSE + GIT_SHALLOW TRUE + ) + endif() + FetchContent_MakeAvailable(cutlass) + + include_directories(${CUTLASS_INCLUDE_DIR}) +else() + message(STATUS "Using system cutlass with version: ${NvidiaCutlass_VERSION}") +endif(NOT NvidiaCutlass_FOUND) + +endif(GPU_LANG STREQUAL "CUDA") +if(GPU_LANG STREQUAL "CUDA") + get_torch_gpu_compiler_flags(TORCH_GPU_FLAGS ${GPU_LANG}) + list(APPEND GPU_FLAGS ${TORCH_GPU_FLAGS}) +endif() + +set(TORCH_natten_SRC + torch-ext/torch_binding.cpp +) + +# TODO: check if CLion support this: +# https://youtrack.jetbrains.com/issue/CPP-16510/CLion-does-not-handle-per-file-include-directories +set_source_files_properties( + ${TORCH_natten_SRC} + PROPERTIES INCLUDE_DIRECTORIES "${CMAKE_SOURCE_DIR}/csrc/include") + + +list(APPEND SRC "${TORCH_natten_SRC}") +if(GPU_LANG STREQUAL "CUDA") +cuda_kernel_component(SRC + NAME natten_blackwell + SOURCES "csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_cm.h" +"csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_dtype.h" +"csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_head_dim.h" +"csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_tile_shape.h" +"csrc/autogen/include/natten_autogen/cuda/blackwell_fna/interface.h" +"csrc/autogen/include/natten_autogen/cuda/blackwell_fna/kernels.h" +"csrc/autogen/src/cuda/blackwell_fna/source_0.cu" +"csrc/autogen/src/cuda/blackwell_fna/source_1.cu" +"csrc/autogen/src/cuda/blackwell_fna/source_10.cu" +"csrc/autogen/src/cuda/blackwell_fna/source_11.cu" +"csrc/autogen/src/cuda/blackwell_fna/source_12.cu" +"csrc/autogen/src/cuda/blackwell_fna/source_13.cu" +"csrc/autogen/src/cuda/blackwell_fna/source_14.cu" +"csrc/autogen/src/cuda/blackwell_fna/source_15.cu" +"csrc/autogen/src/cuda/blackwell_fna/source_16.cu" +"csrc/autogen/src/cuda/blackwell_fna/source_17.cu" +"csrc/autogen/src/cuda/blackwell_fna/source_18.cu" +"csrc/autogen/src/cuda/blackwell_fna/source_19.cu" +"csrc/autogen/src/cuda/blackwell_fna/source_2.cu" +"csrc/autogen/src/cuda/blackwell_fna/source_20.cu" +"csrc/autogen/src/cuda/blackwell_fna/source_21.cu" +"csrc/autogen/src/cuda/blackwell_fna/source_22.cu" +"csrc/autogen/src/cuda/blackwell_fna/source_23.cu" +"csrc/autogen/src/cuda/blackwell_fna/source_24.cu" +"csrc/autogen/src/cuda/blackwell_fna/source_25.cu" +"csrc/autogen/src/cuda/blackwell_fna/source_26.cu" +"csrc/autogen/src/cuda/blackwell_fna/source_27.cu" +"csrc/autogen/src/cuda/blackwell_fna/source_3.cu" +"csrc/autogen/src/cuda/blackwell_fna/source_4.cu" +"csrc/autogen/src/cuda/blackwell_fna/source_5.cu" +"csrc/autogen/src/cuda/blackwell_fna/source_6.cu" +"csrc/autogen/src/cuda/blackwell_fna/source_7.cu" +"csrc/autogen/src/cuda/blackwell_fna/source_8.cu" +"csrc/autogen/src/cuda/blackwell_fna/source_9.cu" +"csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_cm.h" +"csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_dtype.h" +"csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_head_dim.h" +"csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_tile_shape.h" +"csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/interface.h" +"csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/kernels.h" +"csrc/autogen/src/cuda/blackwell_fna_bwd/source_0.cu" +"csrc/autogen/src/cuda/blackwell_fna_bwd/source_1.cu" +"csrc/autogen/src/cuda/blackwell_fna_bwd/source_10.cu" +"csrc/autogen/src/cuda/blackwell_fna_bwd/source_11.cu" +"csrc/autogen/src/cuda/blackwell_fna_bwd/source_12.cu" +"csrc/autogen/src/cuda/blackwell_fna_bwd/source_13.cu" +"csrc/autogen/src/cuda/blackwell_fna_bwd/source_2.cu" +"csrc/autogen/src/cuda/blackwell_fna_bwd/source_3.cu" +"csrc/autogen/src/cuda/blackwell_fna_bwd/source_4.cu" +"csrc/autogen/src/cuda/blackwell_fna_bwd/source_5.cu" +"csrc/autogen/src/cuda/blackwell_fna_bwd/source_6.cu" +"csrc/autogen/src/cuda/blackwell_fna_bwd/source_7.cu" +"csrc/autogen/src/cuda/blackwell_fna_bwd/source_8.cu" +"csrc/autogen/src/cuda/blackwell_fna_bwd/source_9.cu" +"csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/dispatch_head_dim.h" +"csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/dispatch_tile_size.h" +"csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/interface.h" +"csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/kernels.h" +"csrc/autogen/src/cuda/blackwell_fmha/source_0.cu" +"csrc/autogen/src/cuda/blackwell_fmha/source_1.cu" +"csrc/autogen/src/cuda/blackwell_fmha/source_2.cu" +"csrc/autogen/src/cuda/blackwell_fmha/source_3.cu" +"csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/dispatch_head_dim.h" +"csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/dispatch_tile_size.h" +"csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/interface.h" +"csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/kernels.h" +"csrc/autogen/src/cuda/blackwell_fmha_bwd/source_0.cu" +"csrc/autogen/src/cuda/blackwell_fmha_bwd/source_1.cu" +"csrc/autogen/src/cuda/blackwell_fmha_bwd/source_2.cu" +"csrc/autogen/src/cuda/blackwell_fmha_bwd/source_3.cu" + CUDA_MINVER 12.8 INCLUDES "${CMAKE_SOURCE_DIR}/csrc/include;${CMAKE_SOURCE_DIR}/csrc/autogen/include" CUDA_CAPABILITIES 10.0a CUDA_FLAGS "-O3;-std=c++17;--expt-relaxed-constexpr;--extended-lambda;--use_fast_math;--ftemplate-backtrace-limit=0;-Xfatbin=-compress-all;-DNDEBUG;-DNATTEN_WITH_CUTLASS;-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1;-DNATTEN_WITH_BLACKWELL_FNA=1" ) +endif() +if(GPU_LANG STREQUAL "CUDA") +cuda_kernel_component(SRC + NAME natten + SOURCES "csrc/include/natten/blackwell_fmha.h" +"csrc/include/natten/blackwell_fna.h" +"csrc/include/natten/compute_delta.h" +"csrc/include/natten/cuda/fmha/epilogue/epilogue_pipelined.h" +"csrc/include/natten/cuda/fmha/epilogue/epilogue_rescale_output.h" +"csrc/include/natten/cuda/fmha/epilogue/epilogue_thread_apply_logsumexp.h" +"csrc/include/natten/cuda/fmha/fmha_backward.cuh" +"csrc/include/natten/cuda/fmha/fmha_forward.cuh" +"csrc/include/natten/cuda/fmha/gemm/custom_mma.h" +"csrc/include/natten/cuda/fmha/gemm/custom_mma_base.h" +"csrc/include/natten/cuda/fmha/gemm/custom_mma_multistage.h" +"csrc/include/natten/cuda/fmha/gemm/custom_mma_pipelined.h" +"csrc/include/natten/cuda/fmha/gemm/find_default_mma.h" +"csrc/include/natten/cuda/fmha/gemm/mma_accum_lambda_iterator.h" +"csrc/include/natten/cuda/fmha/gemm/mma_from_smem.h" +"csrc/include/natten/cuda/fmha/gemm_kernel_utils.h" +"csrc/include/natten/cuda/fmha/iterators/default_warp_iterator_from_smem.h" +"csrc/include/natten/cuda/fmha/iterators/epilogue_predicated_tile_iterator.h" +"csrc/include/natten/cuda/fmha/iterators/make_residual_last.h" +"csrc/include/natten/cuda/fmha/iterators/predicated_tile_access_iterator_residual_last.h" +"csrc/include/natten/cuda/fmha/iterators/predicated_tile_iterator_residual_last.h" +"csrc/include/natten/cuda/fmha/iterators/transpose_warp_iterator.h" +"csrc/include/natten/cuda/fmha/iterators/warp_iterator_from_smem.h" +"csrc/include/natten/cuda/fmha/kernel_backward.h" +"csrc/include/natten/cuda/fmha/kernel_forward.h" +"csrc/include/natten/cuda/fmha/transform/tile_smem_loader.h" +"csrc/include/natten/cuda/fmha_blackwell/collective/fmha_common.hpp" +"csrc/include/natten/cuda/fmha_blackwell/collective/fmha_fusion.hpp" +"csrc/include/natten/cuda/fmha_blackwell/collective/sm100_fmha_fwd_epilogue_tma_warpspecialized.hpp" +"csrc/include/natten/cuda/fmha_blackwell/collective/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp" +"csrc/include/natten/cuda/fmha_blackwell/collective/sm100_fmha_load_tma_warpspecialized.hpp" +"csrc/include/natten/cuda/fmha_blackwell/common/pow_2.hpp" +"csrc/include/natten/cuda/fmha_blackwell/device/fmha_bwd_sm100.hpp" +"csrc/include/natten/cuda/fmha_blackwell/device/fmha_sm100.hpp" +"csrc/include/natten/cuda/fmha_blackwell/fmha_backward.cuh" +"csrc/include/natten/cuda/fmha_blackwell/fmha_forward.cuh" +"csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_causal_tile_scheduler.hpp" +"csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_kernel_bwd_convert.hpp" +"csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_kernel_bwd_sum_OdO.hpp" +"csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_tile_scheduler.hpp" +"csrc/include/natten/cuda/fmha_blackwell/kernel/sm100_fmha_bwd_kernel_tma_warpspecialized.hpp" +"csrc/include/natten/cuda/fmha_blackwell/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp" +"csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_bwd_tma_warpspecialized.hpp" +"csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_load.hpp" +"csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_softmax.hpp" +"csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_tma.hpp" +"csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_tma_warpspecialized.hpp" +"csrc/include/natten/cuda/fmha_hopper/collective/fmha_common.hpp" +"csrc/include/natten/cuda/fmha_hopper/collective/fmha_epilogue.hpp" +"csrc/include/natten/cuda/fmha_hopper/collective/fmha_epilogue_bwd.hpp" +"csrc/include/natten/cuda/fmha_hopper/collective/fmha_fusion.hpp" +"csrc/include/natten/cuda/fmha_hopper/collective/fmha_varlen.hpp" +"csrc/include/natten/cuda/fmha_hopper/device/fmha_bwd_sm90.hpp" +"csrc/include/natten/cuda/fmha_hopper/device/fmha_sm90.hpp" +"csrc/include/natten/cuda/fmha_hopper/fmha_backward.cuh" +"csrc/include/natten/cuda/fmha_hopper/fmha_forward.cuh" +"csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_builder.hpp" +"csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_bwd_convert.hpp" +"csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_bwd_sum_OdO.hpp" +"csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_tma.hpp" +"csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_tma_warpspecialized.hpp" +"csrc/include/natten/cuda/fmha_hopper/kernel/fmha_options.hpp" +"csrc/include/natten/cuda/fmha_hopper/kernel/fmha_tile_scheduler.hpp" +"csrc/include/natten/cuda/fna/epilogue/epilogue_pipelined.h" +"csrc/include/natten/cuda/fna/epilogue/epilogue_rescale_output.h" +"csrc/include/natten/cuda/fna/epilogue/epilogue_thread_apply_logsumexp.h" +"csrc/include/natten/cuda/fna/epilogue/predicated_tile_iterator.h" +"csrc/include/natten/cuda/fna/epilogue/predicated_tile_iterator_params.h" +"csrc/include/natten/cuda/fna/fna_backward.cuh" +"csrc/include/natten/cuda/fna/fna_forward.cuh" +"csrc/include/natten/cuda/fna/gemm/custom_mma.h" +"csrc/include/natten/cuda/fna/gemm/custom_mma_base.h" +"csrc/include/natten/cuda/fna/gemm/custom_mma_multistage.h" +"csrc/include/natten/cuda/fna/gemm/custom_mma_pipelined.h" +"csrc/include/natten/cuda/fna/gemm/find_default_mma.h" +"csrc/include/natten/cuda/fna/gemm/mma_accum_lambda_iterator.h" +"csrc/include/natten/cuda/fna/gemm/mma_from_smem.h" +"csrc/include/natten/cuda/fna/gemm/replace_mma_iterators.h" +"csrc/include/natten/cuda/fna/gemm_kernel_utils.h" +"csrc/include/natten/cuda/fna/iterators/default_warp_iterator_from_smem.h" +"csrc/include/natten/cuda/fna/iterators/epilogue_predicated_tile_iterator.h" +"csrc/include/natten/cuda/fna/iterators/make_residual_last.h" +"csrc/include/natten/cuda/fna/iterators/predicated_tile_access_iterator.h" +"csrc/include/natten/cuda/fna/iterators/predicated_tile_access_iterator_residual_last.h" +"csrc/include/natten/cuda/fna/iterators/predicated_tile_iterator.h" +"csrc/include/natten/cuda/fna/iterators/predicated_tile_iterator_residual_last.h" +"csrc/include/natten/cuda/fna/iterators/transpose_warp_iterator.h" +"csrc/include/natten/cuda/fna/iterators/warp_iterator_from_smem.h" +"csrc/include/natten/cuda/fna/kernel_backward.h" +"csrc/include/natten/cuda/fna/kernel_forward.h" +"csrc/include/natten/cuda/fna/na_utils.cuh" +"csrc/include/natten/cuda/fna_blackwell/collective/fna_common.hpp" +"csrc/include/natten/cuda/fna_blackwell/collective/fna_fusion.hpp" +"csrc/include/natten/cuda/fna_blackwell/collective/fna_fusion_bwd.hpp" +"csrc/include/natten/cuda/fna_blackwell/collective/sm100_fna_fwd_epilogue_tma_warpspecialized.hpp" +"csrc/include/natten/cuda/fna_blackwell/collective/sm100_fna_fwd_mainloop_tma_warpspecialized.hpp" +"csrc/include/natten/cuda/fna_blackwell/collective/sm100_fna_load_tma_warpspecialized.hpp" +"csrc/include/natten/cuda/fna_blackwell/common/pow_2.hpp" +"csrc/include/natten/cuda/fna_blackwell/device/fna_bwd_sm100.hpp" +"csrc/include/natten/cuda/fna_blackwell/device/fna_sm100.hpp" +"csrc/include/natten/cuda/fna_blackwell/fna_backward.cuh" +"csrc/include/natten/cuda/fna_blackwell/fna_forward.cuh" +"csrc/include/natten/cuda/fna_blackwell/kernel/sm100_fna_bwd_kernel_tma_warpspecialized.hpp" +"csrc/include/natten/cuda/fna_blackwell/kernel/sm100_fna_fwd_kernel_tma_warpspecialized.hpp" +"csrc/include/natten/cuda/fna_hopper/collective/fna_collective_bwd_tma_warpspecialized.hpp" +"csrc/include/natten/cuda/fna_hopper/collective/fna_collective_load.hpp" +"csrc/include/natten/cuda/fna_hopper/collective/fna_collective_softmax.hpp" +"csrc/include/natten/cuda/fna_hopper/collective/fna_collective_tma.hpp" +"csrc/include/natten/cuda/fna_hopper/collective/fna_collective_tma_warpspecialized.hpp" +"csrc/include/natten/cuda/fna_hopper/collective/fna_common.hpp" +"csrc/include/natten/cuda/fna_hopper/collective/fna_fusion.hpp" +"csrc/include/natten/cuda/fna_hopper/collective/fna_fusion_bwd.hpp" +"csrc/include/natten/cuda/fna_hopper/device/fna_bwd_sm90.hpp" +"csrc/include/natten/cuda/fna_hopper/device/fna_sm90.hpp" +"csrc/include/natten/cuda/fna_hopper/fna_backward.cuh" +"csrc/include/natten/cuda/fna_hopper/fna_forward.cuh" +"csrc/include/natten/cuda/fna_hopper/kernel/fna_kernel_builder.hpp" +"csrc/include/natten/cuda/hopper_fmha_fna.h" +"csrc/include/natten/cuda/reduction/compute_delta.cuh" +"csrc/include/natten/cuda/reduction/fmha_kernel_bwd_sum_OdO.hpp" +"csrc/include/natten/cuda/reference/fna_reference_backward.hpp" +"csrc/include/natten/cuda/reference/fna_reference_forward.hpp" +"csrc/include/natten/cuda/reference/mask.hpp" +"csrc/include/natten/cuda/reference/utils.hpp" +"csrc/include/natten/cuda/tokperm/layouts.hpp" +"csrc/include/natten/cuda/tokperm/token_permute_kernel.cuh" +"csrc/include/natten/cuda/tokperm/tokperm.hpp" +"csrc/include/natten/cuda/tokperm/utils/permute.cuh" +"csrc/include/natten/cuda/tokperm/utils/stride.cuh" +"csrc/include/natten/cuda/tokperm/utils/tuple.cuh" +"csrc/include/natten/cuda/utils/cuda.h" +"csrc/include/natten/cuda/utils/cutlass.cuh" +"csrc/include/natten/cuda/utils/generic_cutlass_device.hpp" +"csrc/include/natten/fmha.h" +"csrc/include/natten/fna.h" +"csrc/include/natten/helpers.h" +"csrc/include/natten/hopper_fmha.h" +"csrc/include/natten/hopper_fna.h" +"csrc/include/natten/natten.h" +"csrc/include/natten/reference.h" +"csrc/include/natten/token_permute.h" +"csrc/src/blackwell_fmha.cu" +"csrc/src/blackwell_fna_backward.cu" +"csrc/src/blackwell_fna_forward.cu" +"csrc/src/compute_delta.cu" +"csrc/src/fmha.cu" +"csrc/src/fna_backward.cu" +"csrc/src/fna_forward.cu" +"csrc/src/hopper_fmha.cu" +"csrc/src/hopper_fna_backward.cu" +"csrc/src/hopper_fna_forward.cu" +"csrc/src/reference_backward.cu" +"csrc/src/reference_forward.cu" +"csrc/src/token_permute.cu" +"csrc/src/token_unpermute.cu" +"csrc/autogen/include/natten_autogen/cuda/reference/dispatch_cm.h" +"csrc/autogen/include/natten_autogen/cuda/reference/dispatch_dtype.h" +"csrc/autogen/include/natten_autogen/cuda/reference/interface.h" +"csrc/autogen/include/natten_autogen/cuda/reference/kernels.h" +"csrc/autogen/src/cuda/reference/source_0.cu" +"csrc/autogen/src/cuda/reference/source_1.cu" +"csrc/autogen/include/natten_autogen/cuda/fna/dispatch_cm.h" +"csrc/autogen/include/natten_autogen/cuda/fna/dispatch_device.h" +"csrc/autogen/include/natten_autogen/cuda/fna/dispatch_dtype.h" +"csrc/autogen/include/natten_autogen/cuda/fna/interface.h" +"csrc/autogen/include/natten_autogen/cuda/fna/kernels.h" +"csrc/autogen/src/cuda/fna/source_0.cu" +"csrc/autogen/src/cuda/fna/source_1.cu" +"csrc/autogen/src/cuda/fna/source_10.cu" +"csrc/autogen/src/cuda/fna/source_11.cu" +"csrc/autogen/src/cuda/fna/source_12.cu" +"csrc/autogen/src/cuda/fna/source_13.cu" +"csrc/autogen/src/cuda/fna/source_14.cu" +"csrc/autogen/src/cuda/fna/source_15.cu" +"csrc/autogen/src/cuda/fna/source_16.cu" +"csrc/autogen/src/cuda/fna/source_17.cu" +"csrc/autogen/src/cuda/fna/source_18.cu" +"csrc/autogen/src/cuda/fna/source_19.cu" +"csrc/autogen/src/cuda/fna/source_2.cu" +"csrc/autogen/src/cuda/fna/source_20.cu" +"csrc/autogen/src/cuda/fna/source_21.cu" +"csrc/autogen/src/cuda/fna/source_22.cu" +"csrc/autogen/src/cuda/fna/source_23.cu" +"csrc/autogen/src/cuda/fna/source_24.cu" +"csrc/autogen/src/cuda/fna/source_25.cu" +"csrc/autogen/src/cuda/fna/source_26.cu" +"csrc/autogen/src/cuda/fna/source_27.cu" +"csrc/autogen/src/cuda/fna/source_28.cu" +"csrc/autogen/src/cuda/fna/source_29.cu" +"csrc/autogen/src/cuda/fna/source_3.cu" +"csrc/autogen/src/cuda/fna/source_30.cu" +"csrc/autogen/src/cuda/fna/source_31.cu" +"csrc/autogen/src/cuda/fna/source_32.cu" +"csrc/autogen/src/cuda/fna/source_33.cu" +"csrc/autogen/src/cuda/fna/source_34.cu" +"csrc/autogen/src/cuda/fna/source_35.cu" +"csrc/autogen/src/cuda/fna/source_36.cu" +"csrc/autogen/src/cuda/fna/source_37.cu" +"csrc/autogen/src/cuda/fna/source_38.cu" +"csrc/autogen/src/cuda/fna/source_39.cu" +"csrc/autogen/src/cuda/fna/source_4.cu" +"csrc/autogen/src/cuda/fna/source_40.cu" +"csrc/autogen/src/cuda/fna/source_41.cu" +"csrc/autogen/src/cuda/fna/source_42.cu" +"csrc/autogen/src/cuda/fna/source_43.cu" +"csrc/autogen/src/cuda/fna/source_44.cu" +"csrc/autogen/src/cuda/fna/source_45.cu" +"csrc/autogen/src/cuda/fna/source_46.cu" +"csrc/autogen/src/cuda/fna/source_47.cu" +"csrc/autogen/src/cuda/fna/source_48.cu" +"csrc/autogen/src/cuda/fna/source_49.cu" +"csrc/autogen/src/cuda/fna/source_5.cu" +"csrc/autogen/src/cuda/fna/source_50.cu" +"csrc/autogen/src/cuda/fna/source_51.cu" +"csrc/autogen/src/cuda/fna/source_52.cu" +"csrc/autogen/src/cuda/fna/source_53.cu" +"csrc/autogen/src/cuda/fna/source_54.cu" +"csrc/autogen/src/cuda/fna/source_55.cu" +"csrc/autogen/src/cuda/fna/source_56.cu" +"csrc/autogen/src/cuda/fna/source_57.cu" +"csrc/autogen/src/cuda/fna/source_58.cu" +"csrc/autogen/src/cuda/fna/source_59.cu" +"csrc/autogen/src/cuda/fna/source_6.cu" +"csrc/autogen/src/cuda/fna/source_60.cu" +"csrc/autogen/src/cuda/fna/source_61.cu" +"csrc/autogen/src/cuda/fna/source_62.cu" +"csrc/autogen/src/cuda/fna/source_63.cu" +"csrc/autogen/src/cuda/fna/source_7.cu" +"csrc/autogen/src/cuda/fna/source_8.cu" +"csrc/autogen/src/cuda/fna/source_9.cu" +"csrc/autogen/include/natten_autogen/cuda/fmha/dispatch_dtype.h" +"csrc/autogen/include/natten_autogen/cuda/fmha/interface.h" +"csrc/autogen/include/natten_autogen/cuda/fmha/kernels.h" +"csrc/autogen/src/cuda/fmha/source_0.cu" +"csrc/autogen/src/cuda/fmha/source_1.cu" +"csrc/autogen/src/cuda/fmha/source_2.cu" +"csrc/autogen/src/cuda/fmha/source_3.cu" +"csrc/autogen/src/cuda/fmha/source_4.cu" +"csrc/autogen/src/cuda/fmha/source_5.cu" + INCLUDES "${CMAKE_SOURCE_DIR}/csrc/include;${CMAKE_SOURCE_DIR}/csrc/autogen/include" CUDA_CAPABILITIES 7.0 7.5 8.0 8.6 8.9 9.0 10.0 12.0 CUDA_FLAGS "-O3;-std=c++17;--expt-relaxed-constexpr;--extended-lambda;--use_fast_math;--ftemplate-backtrace-limit=0;-Xfatbin=-compress-all;-DNDEBUG;-DNATTEN_WITH_CUTLASS;-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1" ) +endif() +if(GPU_LANG STREQUAL "CUDA") +cuda_kernel_component(SRC + NAME natten_hopper + SOURCES "csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_cm.h" +"csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_dtype.h" +"csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_head_dim.h" +"csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_tile_shape.h" +"csrc/autogen/include/natten_autogen/cuda/hopper_fna/interface.h" +"csrc/autogen/include/natten_autogen/cuda/hopper_fna/kernels.h" +"csrc/autogen/src/cuda/hopper_fna/source_0.cu" +"csrc/autogen/src/cuda/hopper_fna/source_1.cu" +"csrc/autogen/src/cuda/hopper_fna/source_2.cu" +"csrc/autogen/src/cuda/hopper_fna/source_3.cu" +"csrc/autogen/src/cuda/hopper_fna/source_4.cu" +"csrc/autogen/src/cuda/hopper_fna/source_5.cu" +"csrc/autogen/src/cuda/hopper_fna/source_6.cu" +"csrc/autogen/src/cuda/hopper_fna/source_7.cu" +"csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_cm.h" +"csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_dtype.h" +"csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_head_dim.h" +"csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_tile_shape.h" +"csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/interface.h" +"csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/kernels.h" +"csrc/autogen/src/cuda/hopper_fna_bwd/source_0.cu" +"csrc/autogen/src/cuda/hopper_fna_bwd/source_1.cu" +"csrc/autogen/src/cuda/hopper_fna_bwd/source_2.cu" +"csrc/autogen/src/cuda/hopper_fna_bwd/source_3.cu" +"csrc/autogen/include/natten_autogen/cuda/hopper_fmha/dispatch_head_dim.h" +"csrc/autogen/include/natten_autogen/cuda/hopper_fmha/dispatch_tile_size.h" +"csrc/autogen/include/natten_autogen/cuda/hopper_fmha/interface.h" +"csrc/autogen/include/natten_autogen/cuda/hopper_fmha/kernels.h" +"csrc/autogen/src/cuda/hopper_fmha/source_0.cu" +"csrc/autogen/src/cuda/hopper_fmha/source_1.cu" +"csrc/autogen/src/cuda/hopper_fmha/source_2.cu" +"csrc/autogen/src/cuda/hopper_fmha/source_3.cu" +"csrc/autogen/src/cuda/hopper_fmha/source_4.cu" +"csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/dispatch_head_dim.h" +"csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/dispatch_tile_size.h" +"csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/interface.h" +"csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/kernels.h" +"csrc/autogen/src/cuda/hopper_fmha_bwd/source_0.cu" +"csrc/autogen/src/cuda/hopper_fmha_bwd/source_1.cu" +"csrc/autogen/src/cuda/hopper_fmha_bwd/source_2.cu" +"csrc/autogen/src/cuda/hopper_fmha_bwd/source_3.cu" +"csrc/autogen/src/cuda/hopper_fmha_bwd/source_4.cu" + INCLUDES "${CMAKE_SOURCE_DIR}/csrc/include;${CMAKE_SOURCE_DIR}/csrc/autogen/include" CUDA_CAPABILITIES 9.0a CUDA_FLAGS "-O3;-std=c++17;--expt-relaxed-constexpr;--extended-lambda;--use_fast_math;--ftemplate-backtrace-limit=0;-Xfatbin=-compress-all;-DNDEBUG;-DNATTEN_WITH_CUTLASS;-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1;-DNATTEN_WITH_HOPPER_FNA=1" ) +endif() +# Include Metal shader compilation utilities if needed +if(GPU_LANG STREQUAL "METAL") + include(${CMAKE_CURRENT_LIST_DIR}/cmake/compile-metal.cmake) +endif() + +# Define the extension target with unified parameters +define_gpu_extension_target( + ${OPS_NAME} + ${OPS_NAME} + DESTINATION ${OPS_NAME} + LANGUAGE ${GPU_LANG} + SOURCES ${SRC} + COMPILE_FLAGS ${GPU_FLAGS} + USE_SABI 3 + WITH_SOABI) + +if(GPU_LANG STREQUAL "SYCL") + target_link_options(${OPS_NAME} PRIVATE ${sycl_link_flags}) + target_link_libraries(${OPS_NAME} PRIVATE dnnl) +endif() + +# Compile Metal shaders if any were found +if(GPU_LANG STREQUAL "METAL") + if(ALL_METAL_SOURCES) + compile_metal_shaders(${OPS_NAME} "${ALL_METAL_SOURCES}" "${METAL_INCLUDE_DIRS}") + endif() +endif() + + +# Add kernels_install target for huggingface/kernels library layout +add_kernels_install_target(${OPS_NAME} "natten" "${BUILD_VARIANT_NAME}" + DATA_EXTENSIONS "" + GPU_ARCHS "${ALL_GPU_ARCHS}") + +# Add local_install target for local development with get_local_kernel() +add_local_install_target(${OPS_NAME} "natten" "${BUILD_VARIANT_NAME}" + DATA_EXTENSIONS "" + GPU_ARCHS "${ALL_GPU_ARCHS}") diff --git a/natten/cmake/_ops.py.in b/natten/cmake/_ops.py.in new file mode 100644 index 00000000..736771e8 --- /dev/null +++ b/natten/cmake/_ops.py.in @@ -0,0 +1,9 @@ +import torch +from . import @OPS_NAME@ +ops = torch.ops.@OPS_NAME@ + +def add_op_namespace_prefix(op_name: str): + """ + Prefix op by namespace. + """ + return f"@OPS_NAME@::{op_name}" diff --git a/natten/cmake/add_build_metadata.py b/natten/cmake/add_build_metadata.py new file mode 100644 index 00000000..23a02492 --- /dev/null +++ b/natten/cmake/add_build_metadata.py @@ -0,0 +1,56 @@ +import argparse +import json +import sys + + +def main(): + parser = argparse.ArgumentParser( + description="Write a metadata JSON file with GPU architecture information, " + "reading from a source file and writing to a destination." + ) + parser.add_argument( + "input", + help="Path to the source metadata JSON file to read from.", + ) + parser.add_argument( + "destination", + help="Path to write the output metadata JSON file to.", + ) + + parser.add_argument( + "--archs", + help="Semicolon-separated list of GPU architectures/capabilities.", + ) + + args = parser.parse_args() + + archs = ( + sorted(set(a for a in args.archs.split(";") if a)) + if args.archs is not None + else None + ) + + try: + with open(args.input) as f: + data = json.load(f) + except FileNotFoundError: + print(f"Error: input metadata file not found: {args.input}", file=sys.stderr) + sys.exit(1) + except json.JSONDecodeError as e: + print(f"Error: failed to parse input metadata JSON: {e}", file=sys.stderr) + sys.exit(1) + + if archs is not None: + data["backend"]["archs"] = archs + + try: + with open(args.destination, "w") as f: + json.dump(data, f, indent=2) + f.write("\n") + except OSError as e: + print(f"Error: failed to write output metadata JSON: {e}", file=sys.stderr) + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/natten/cmake/build-variants.cmake b/natten/cmake/build-variants.cmake new file mode 100644 index 00000000..31fb8526 --- /dev/null +++ b/natten/cmake/build-variants.cmake @@ -0,0 +1,328 @@ +# Generate a standardized build variant name following the pattern: +# torch-[cxx11-]-- +# or, when compiled against the Torch stable ABI: +# torch-stable-abi--- +# +# Arguments: +# OUT_BUILD_NAME - Output variable name +# TORCH_VERSION - PyTorch version (e.g., "2.7.1"); ignored when TORCH_STABLE_ABI is set +# COMPUTE_FRAMEWORK - One of: cuda, rocm, metal, xpu, cpu +# COMPUTE_VERSION - Version of compute framework (e.g., "12.4" for CUDA, "6.0" for ROCm) +# Optional for CPU-only builds (pass empty string or omit) +# Optional keyword arguments: +# TORCH_STABLE_ABI - Stable ABI version the extension was compiled against (e.g., "2.11"); +# when set, TORCH_VERSION is ignored and the prefix becomes +# torch-stable-abi (e.g., "2.11" -> "torch-stable-abi211") +# Example output: torch27-cxx11-cu124-x86_64-linux (Linux) +# torch27-cu124-x86_64-windows (Windows) +# torch27-metal-aarch64-darwin (macOS) +# torch-stable-abi211-cu124-x86_64-linux (Linux, stable ABI) +# +function(generate_build_name OUT_BUILD_NAME TORCH_VERSION COMPUTE_FRAMEWORK COMPUTE_VERSION) + cmake_parse_arguments(ARG "" "TORCH_STABLE_ABI" "" ${ARGN}) + + if(ARG_TORCH_STABLE_ABI) + # Flatten the stable ABI version (e.g., "2.11" -> "211") and ignore TORCH_VERSION + string(REPLACE "." ";" VERSION_LIST "${ARG_TORCH_STABLE_ABI}") + list(LENGTH VERSION_LIST VERSION_COMPONENTS) + if(VERSION_COMPONENTS LESS 2) + list(APPEND VERSION_LIST "0") + endif() + list(GET VERSION_LIST 0 MAJOR) + list(GET VERSION_LIST 1 MINOR) + set(TORCH_PREFIX "torch-stable-abi${MAJOR}${MINOR}") + else() + # Flatten version by removing dots and padding to 2 components + string(REPLACE "." ";" VERSION_LIST "${TORCH_VERSION}") + list(LENGTH VERSION_LIST VERSION_COMPONENTS) + + # Pad to at least 2 components + if(VERSION_COMPONENTS LESS 2) + list(APPEND VERSION_LIST "0") + endif() + + # Take first 2 components and join without dots + list(GET VERSION_LIST 0 MAJOR) + list(GET VERSION_LIST 1 MINOR) + set(TORCH_PREFIX "torch${MAJOR}${MINOR}") + endif() + + # Generate compute string + if(COMPUTE_FRAMEWORK STREQUAL "cuda") + # Flatten CUDA version (e.g., "12.4" -> "124") + string(REPLACE "." ";" COMPUTE_VERSION_LIST "${COMPUTE_VERSION}") + list(LENGTH COMPUTE_VERSION_LIST COMPUTE_COMPONENTS) + if(COMPUTE_COMPONENTS GREATER_EQUAL 2) + list(GET COMPUTE_VERSION_LIST 0 COMPUTE_MAJOR) + list(GET COMPUTE_VERSION_LIST 1 COMPUTE_MINOR) + set(COMPUTE_STRING "cu${COMPUTE_MAJOR}${COMPUTE_MINOR}") + else() + list(GET COMPUTE_VERSION_LIST 0 COMPUTE_MAJOR) + set(COMPUTE_STRING "cu${COMPUTE_MAJOR}0") + endif() + elseif(COMPUTE_FRAMEWORK STREQUAL "rocm") + # Flatten ROCm version (e.g., "6.0" -> "60") + string(REPLACE "." ";" COMPUTE_VERSION_LIST "${COMPUTE_VERSION}") + list(LENGTH COMPUTE_VERSION_LIST COMPUTE_COMPONENTS) + if(COMPUTE_COMPONENTS GREATER_EQUAL 2) + list(GET COMPUTE_VERSION_LIST 0 COMPUTE_MAJOR) + list(GET COMPUTE_VERSION_LIST 1 COMPUTE_MINOR) + set(COMPUTE_STRING "rocm${COMPUTE_MAJOR}${COMPUTE_MINOR}") + else() + list(GET COMPUTE_VERSION_LIST 0 COMPUTE_MAJOR) + set(COMPUTE_STRING "rocm${COMPUTE_MAJOR}0") + endif() + elseif(COMPUTE_FRAMEWORK STREQUAL "xpu") + # Flatten XPU version (e.g., "2025.2" -> "202552") + string(REPLACE "." ";" COMPUTE_VERSION_LIST "${COMPUTE_VERSION}") + list(LENGTH COMPUTE_VERSION_LIST COMPUTE_COMPONENTS) + if(COMPUTE_COMPONENTS GREATER_EQUAL 2) + list(GET COMPUTE_VERSION_LIST 0 COMPUTE_MAJOR) + list(GET COMPUTE_VERSION_LIST 1 COMPUTE_MINOR) + set(COMPUTE_STRING "xpu${COMPUTE_MAJOR}${COMPUTE_MINOR}") + else() + list(GET COMPUTE_VERSION_LIST 0 COMPUTE_MAJOR) + set(COMPUTE_STRING "xpu${COMPUTE_MAJOR}0") + endif() + elseif(COMPUTE_FRAMEWORK STREQUAL "metal") + set(COMPUTE_STRING "metal") + elseif(COMPUTE_FRAMEWORK STREQUAL "cpu") + set(COMPUTE_STRING "cpu") + else() + message(FATAL_ERROR "Unknown compute framework: ${COMPUTE_FRAMEWORK}") + endif() + + # Detect from target system (CMAKE_SYSTEM_* variables refer to target, not host) + # Normalize architecture name + if(CMAKE_SYSTEM_PROCESSOR MATCHES "^(x86_64|amd64|AMD64)$") + set(CPU_ARCH "x86_64") + elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "^(aarch64|arm64|ARM64)$") + set(CPU_ARCH "aarch64") + else() + message(FATAL_ERROR "Unsupported architecture: ${CMAKE_SYSTEM_PROCESSOR}") + endif() + + # Normalize OS name + if(CMAKE_SYSTEM_NAME STREQUAL "Windows") + set(OS_NAME "windows") + elseif(CMAKE_SYSTEM_NAME STREQUAL "Linux") + set(OS_NAME "linux") + elseif(CMAKE_SYSTEM_NAME STREQUAL "Darwin") + set(OS_NAME "darwin") + else() + message(WARNING "Unknown OS ${CMAKE_SYSTEM_NAME}, using as-is") + string(TOLOWER "${CMAKE_SYSTEM_NAME}" OS_NAME) + endif() + + set(ARCH_OS_STRING "${CPU_ARCH}-${OS_NAME}") + + # Assemble the final build name + # For non-stable-ABI Linux builds, include cxx11 ABI indicator for compatibility + if(NOT ARG_TORCH_STABLE_ABI AND ARCH_OS_STRING MATCHES "-linux$") + set(BUILD_NAME "${TORCH_PREFIX}-cxx11-${COMPUTE_STRING}-${ARCH_OS_STRING}") + else() + set(BUILD_NAME "${TORCH_PREFIX}-${COMPUTE_STRING}-${ARCH_OS_STRING}") + endif() + + set(${OUT_BUILD_NAME} "${BUILD_NAME}" PARENT_SCOPE) + message(STATUS "Generated build name: ${BUILD_NAME}") +endfunction() + +# +# Create a custom install target for the huggingface/kernels library layout. +# This installs the extension into a directory structure suitable for kernel hub discovery: +# / +# +# Arguments: +# TARGET_NAME - Name of the target to create the install rule for +# PACKAGE_NAME - Python package name (e.g., "activation") +# BUILD_VARIANT_NAME - Build variant name (e.g., "torch271-cxx11-cu124-x86_64-linux") +# INSTALL_PREFIX - Base installation directory (defaults to CMAKE_INSTALL_PREFIX) +# GPU_ARCHS - List of GPU architectures that were compiled +# (optional; when provided for CUDA/ROCm, metadata.json will include +# a "backend" key with the type and arch list) +# +function(add_kernels_install_target TARGET_NAME PACKAGE_NAME BUILD_VARIANT_NAME) + set(oneValueArgs INSTALL_PREFIX) + set(multiValueArgs DATA_EXTENSIONS GPU_ARCHS) + cmake_parse_arguments(ARG "" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) + + if(NOT ARG_INSTALL_PREFIX) + set(ARG_INSTALL_PREFIX "${CMAKE_INSTALL_PREFIX}") + endif() + + gpu_lang_to_backend(_BACKEND ${GPU_LANG}) + + # Always include 'py' extension for Python files + set(ALL_EXTENSIONS ${ARG_DATA_EXTENSIONS}) + list(APPEND ALL_EXTENSIONS "py") + + # Set the installation directory + set(KERNEL_INSTALL_DIR "${ARG_INSTALL_PREFIX}/${BUILD_VARIANT_NAME}") + + message(STATUS "Using PACKAGE_NAME: ${PACKAGE_NAME}") + + # Install the compiled extension using CMake's install() command + # This will be triggered by the standard INSTALL target + install(TARGETS ${TARGET_NAME} + LIBRARY DESTINATION "${KERNEL_INSTALL_DIR}" + RUNTIME DESTINATION "${KERNEL_INSTALL_DIR}" + COMPONENT ${TARGET_NAME}) + + # Install data files with specified extensions + foreach(ext IN LISTS ALL_EXTENSIONS) + file(GLOB_RECURSE DATA_FILES RELATIVE "${CMAKE_SOURCE_DIR}/torch-ext/${PACKAGE_NAME}" "${CMAKE_SOURCE_DIR}/torch-ext/${PACKAGE_NAME}/*.${ext}") + foreach(data_file IN LISTS DATA_FILES) + get_filename_component(data_file_dir "${data_file}" DIRECTORY) + install(FILES "${CMAKE_SOURCE_DIR}/torch-ext/${PACKAGE_NAME}/${data_file}" + DESTINATION "${KERNEL_INSTALL_DIR}/${data_file_dir}" + COMPONENT ${TARGET_NAME}) + endforeach() + endforeach() + + message(STATUS "GPU archs: ${ARG_GPU_ARCHS}") + + # Add the GPU archs to matadata.json when applicable. + if((GPU_LANG STREQUAL "CUDA" OR GPU_LANG STREQUAL "HIP") AND ARG_GPU_ARCHS) + list(JOIN ARG_GPU_ARCHS ";" _GPU_ARCHS_STR) + install(CODE " + file(MAKE_DIRECTORY \"${KERNEL_INSTALL_DIR}\") + execute_process( + COMMAND \"${Python3_EXECUTABLE}\" + \"${CMAKE_CURRENT_LIST_DIR}/cmake/add_build_metadata.py\" + \"${CMAKE_SOURCE_DIR}/metadata-${_BACKEND}.json\" + \"${KERNEL_INSTALL_DIR}/metadata.json\" + --archs \"${_GPU_ARCHS_STR}\" + RESULT_VARIABLE _METADATA_RESULT + ERROR_VARIABLE _METADATA_ERROR + ) + if(NOT _METADATA_RESULT EQUAL 0) + message(WARNING \"Failed to add GPU arch metadata: \${_METADATA_ERROR}\") + endif() + " COMPONENT ${TARGET_NAME}) + else() + install(CODE " + file(MAKE_DIRECTORY \"${KERNEL_INSTALL_DIR}\") + execute_process( + COMMAND \"${Python3_EXECUTABLE}\" + \"${CMAKE_CURRENT_LIST_DIR}/cmake/add_build_metadata.py\" + \"${CMAKE_SOURCE_DIR}/metadata-${_BACKEND}.json\" + \"${KERNEL_INSTALL_DIR}/metadata.json\" + RESULT_VARIABLE _METADATA_RESULT + ERROR_VARIABLE _METADATA_ERROR + ) + if(NOT _METADATA_RESULT EQUAL 0) + message(WARNING \"Failed to write metadata: \${_METADATA_ERROR}\") + endif() + " COMPONENT ${TARGET_NAME}) + endif() + + # Compatibility with older kernels and direct Python imports. + install(FILES ${CMAKE_SOURCE_DIR}/compat.py + DESTINATION "${KERNEL_INSTALL_DIR}/${PACKAGE_NAME}" + RENAME "__init__.py" + COMPONENT ${TARGET_NAME}) + + message(STATUS "Added install rules for ${TARGET_NAME} -> ${BUILD_VARIANT_NAME}") +endfunction() + +# +# Add install rules for local development with huggingface/kernels. +# This installs the extension into the layout expected by get_local_kernel(): +# ${CMAKE_SOURCE_DIR}/build// +# +# This allows developers to use get_local_kernel() from the kernels library to load +# locally built kernels without needing to publish to the hub. +# +# This uses the standard CMake install() command, so it works with the default +# "install" target that is always available. +# +# Arguments: +# TARGET_NAME - Name of the target to create the install rule for +# PACKAGE_NAME - Python package name (e.g., "activation") +# BUILD_VARIANT_NAME - Build variant name (e.g., "torch271-cxx11-cu124-x86_64-linux") +# GPU_ARCHS - List of GPU architectures that were compiled +# (optional; when provided for CUDA/ROCm, metadata.json will include +# a "backend" key with the type and arch list) +# +function(add_local_install_target TARGET_NAME PACKAGE_NAME BUILD_VARIANT_NAME) + set(oneValueArgs) + set(multiValueArgs DATA_EXTENSIONS GPU_ARCHS) + cmake_parse_arguments(ARG "" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) + + # Always include 'py' extension for Python files + set(ALL_EXTENSIONS ${ARG_DATA_EXTENSIONS}) + list(APPEND ALL_EXTENSIONS "py") + + # Define your local, folder based, installation directory + set(LOCAL_INSTALL_DIR "${CMAKE_SOURCE_DIR}/build/${BUILD_VARIANT_NAME}") + # Variant directory is where metadata.json should go (for kernels upload discovery) + set(VARIANT_DIR "${CMAKE_SOURCE_DIR}/build/${BUILD_VARIANT_NAME}") + + # Create a custom target for local installation + add_custom_target(local_install + COMMENT "Installing files to local directory..." + ) + + gpu_lang_to_backend(_BACKEND ${GPU_LANG}) + + # Copy data files with specified extensions + foreach(ext IN LISTS ALL_EXTENSIONS) + file(GLOB_RECURSE DATA_FILES RELATIVE "${CMAKE_SOURCE_DIR}/torch-ext/${PACKAGE_NAME}" "${CMAKE_SOURCE_DIR}/torch-ext/${PACKAGE_NAME}/*.${ext}") + foreach(data_file IN LISTS DATA_FILES) + get_filename_component(data_file_dir "${data_file}" DIRECTORY) + add_custom_command(TARGET local_install POST_BUILD + COMMAND ${CMAKE_COMMAND} -E make_directory + ${LOCAL_INSTALL_DIR}/${data_file_dir} + COMMAND ${CMAKE_COMMAND} -E copy_if_different + ${CMAKE_SOURCE_DIR}/torch-ext/${PACKAGE_NAME}/${data_file} + ${LOCAL_INSTALL_DIR}/${data_file_dir}/ + COMMENT "Copying ${data_file} to ${LOCAL_INSTALL_DIR}/${data_file_dir}" + ) + endforeach() + endforeach() + + + # Add the GPU archs to matadata.json when applicable. + if((GPU_LANG STREQUAL "CUDA" OR GPU_LANG STREQUAL "HIP") AND ARG_GPU_ARCHS) + list(JOIN ARG_GPU_ARCHS ";" _GPU_ARCHS_STR) + add_custom_command(TARGET local_install POST_BUILD + COMMAND ${CMAKE_COMMAND} -E make_directory ${VARIANT_DIR} + COMMAND ${Python3_EXECUTABLE} + ${CMAKE_CURRENT_LIST_DIR}/cmake/add_build_metadata.py + ${CMAKE_SOURCE_DIR}/metadata-${_BACKEND}.json + ${VARIANT_DIR}/metadata.json + --archs "${_GPU_ARCHS_STR}" + COMMENT "Writing metadata.json with GPU arch info to ${VARIANT_DIR}" + ) + else() + add_custom_command(TARGET local_install POST_BUILD + COMMAND ${CMAKE_COMMAND} -E make_directory ${VARIANT_DIR} + COMMAND ${Python3_EXECUTABLE} + ${CMAKE_CURRENT_LIST_DIR}/cmake/add_build_metadata.py + ${CMAKE_SOURCE_DIR}/metadata-${_BACKEND}.json + ${VARIANT_DIR}/metadata.json + COMMENT "Writing metadata.json to ${VARIANT_DIR}" + ) + endif() + + add_custom_command(TARGET local_install POST_BUILD + # Copy the shared library + COMMAND ${CMAKE_COMMAND} -E copy_if_different + $ + ${LOCAL_INSTALL_DIR}/ + + # Compatibility with older kernels and direct Python imports. + COMMAND ${CMAKE_COMMAND} -E copy_if_different + ${CMAKE_SOURCE_DIR}/compat.py + ${VARIANT_DIR}/${PACKAGE_NAME}/__init__.py + + COMMENT "Copying shared library and Python files to ${LOCAL_INSTALL_DIR}" + COMMAND_EXPAND_LISTS + ) + + # Create both directories: variant dir for metadata.json, package dir for binaries + file(MAKE_DIRECTORY ${VARIANT_DIR}) + file(MAKE_DIRECTORY ${LOCAL_INSTALL_DIR}) + message(STATUS "Added install rules for ${TARGET_NAME} -> build/${BUILD_VARIANT_NAME}") +endfunction() diff --git a/natten/cmake/compile-metal.cmake b/natten/cmake/compile-metal.cmake new file mode 100644 index 00000000..50d44a2d --- /dev/null +++ b/natten/cmake/compile-metal.cmake @@ -0,0 +1,104 @@ +# Metal shader compilation function +function(compile_metal_shaders TARGET_NAME METAL_SOURCES EXTRA_INCLUDE_DIRS) + if(NOT DEFINED METAL_TOOLCHAIN) + execute_process( + COMMAND "xcodebuild" "-showComponent" "MetalToolchain" + OUTPUT_VARIABLE FIND_METAL_OUT + RESULT_VARIABLE FIND_METAL_ERROR_CODE + ERROR_VARIABLE FIND_METAL_STDERR + OUTPUT_STRIP_TRAILING_WHITESPACE) + + if(NOT FIND_METAL_ERROR_CODE EQUAL 0) + message(FATAL_ERROR "${ERR_MSG}: ${FIND_METAL_STDERR}") + endif() + + # Extract the Toolchain Search Path value and append Metal.xctoolchain + string(REGEX MATCH "Toolchain Search Path: ([^\n]+)" MATCH_RESULT "${FIND_METAL_OUT}") + set(METAL_TOOLCHAIN "${CMAKE_MATCH_1}/Metal.xctoolchain") + endif() + + # Set Metal compiler flags + set(METAL_FLAGS "-std=metal4.0" "-O2") + + # Output directory for compiled metallib + set(METALLIB_OUTPUT_DIR "${CMAKE_BINARY_DIR}/metallib") + file(MAKE_DIRECTORY ${METALLIB_OUTPUT_DIR}) + + foreach(INC ${EXTRA_INCLUDE_DIRS}) + list(APPEND METAL_FLAGS "-I${INC}") + endforeach() + + # Separate .metal files from .h files and compile .metal files to .air + set(AIR_FILES) + set(METAL_FILES) + set(HEADER_FILES) + + foreach(SOURCE_FILE ${METAL_SOURCES}) + if(SOURCE_FILE MATCHES "\\.metal$") + list(APPEND METAL_FILES ${SOURCE_FILE}) + elseif(SOURCE_FILE MATCHES "\\.h$") + list(APPEND HEADER_FILES ${SOURCE_FILE}) + endif() + endforeach() + + foreach(METAL_FILE ${METAL_FILES}) + get_filename_component(METAL_NAME ${METAL_FILE} NAME_WE) + set(AIR_FILE "${CMAKE_BINARY_DIR}/${METAL_NAME}.air") + + # Include header files as dependencies + set(ALL_DEPENDENCIES ${CMAKE_CURRENT_SOURCE_DIR}/${METAL_FILE}) + foreach(HEADER_FILE ${HEADER_FILES}) + list(APPEND ALL_DEPENDENCIES ${CMAKE_CURRENT_SOURCE_DIR}/${HEADER_FILE}) + endforeach() + + add_custom_command( + OUTPUT ${AIR_FILE} + COMMAND "${METAL_TOOLCHAIN}/usr/bin/metal" ${METAL_FLAGS} + -c ${CMAKE_CURRENT_SOURCE_DIR}/${METAL_FILE} + -o ${AIR_FILE} + DEPENDS ${ALL_DEPENDENCIES} + COMMENT "Compiling Metal shader ${METAL_FILE} to ${AIR_FILE}" + VERBATIM + ) + + list(APPEND AIR_FILES ${AIR_FILE}) + endforeach() + + # Link all .air files into a single .metallib + set(METALLIB_FILE "${METALLIB_OUTPUT_DIR}/${TARGET_NAME}.metallib") + add_custom_command( + OUTPUT ${METALLIB_FILE} + COMMAND "${METAL_TOOLCHAIN}/usr/bin/metallib" ${AIR_FILES} + -o ${METALLIB_FILE} + DEPENDS ${AIR_FILES} + COMMENT "Linking Metal library ${METALLIB_FILE}" + VERBATIM + ) + + # Generate C++ header with embedded metallib data + set(METALLIB_HEADER "${CMAKE_BINARY_DIR}/${TARGET_NAME}_metallib.h") + set(METALLIB_TO_HEADER_SCRIPT "${CMAKE_CURRENT_SOURCE_DIR}/cmake/metallib_to_header.py") + + add_custom_command( + OUTPUT ${METALLIB_HEADER} + COMMAND ${Python3_EXECUTABLE} ${METALLIB_TO_HEADER_SCRIPT} ${METALLIB_FILE} ${METALLIB_HEADER} ${TARGET_NAME} + DEPENDS ${METALLIB_FILE} ${METALLIB_TO_HEADER_SCRIPT} + COMMENT "Generating embedded Metal library header ${METALLIB_HEADER}" + VERBATIM + ) + + # Create a custom target for the metallib + add_custom_target(${TARGET_NAME}_metallib ALL DEPENDS ${METALLIB_FILE} ${METALLIB_HEADER}) + + # Add dependency to main target + add_dependencies(${TARGET_NAME} ${TARGET_NAME}_metallib) + + # Add the generated header to include directories + target_include_directories(${TARGET_NAME} PRIVATE ${CMAKE_BINARY_DIR}) + + # Pass the metallib header and namespace as compile definitions + target_compile_definitions(${TARGET_NAME} PRIVATE + EMBEDDED_METALLIB_HEADER="${TARGET_NAME}_metallib.h" + EMBEDDED_METALLIB_NAMESPACE=${TARGET_NAME}_metal + ) +endfunction() diff --git a/natten/cmake/get_gpu_lang.cmake b/natten/cmake/get_gpu_lang.cmake new file mode 100644 index 00000000..277ac7b0 --- /dev/null +++ b/natten/cmake/get_gpu_lang.cmake @@ -0,0 +1,9 @@ +# +# Get the GPU language from Torch. +# +function(get_gpu_lang OUT) + run_python_script(PYTHON_OUT + "${CMAKE_CURRENT_SOURCE_DIR}/cmake/get_gpu_lang.py" + "Cannot detect GPU language") + set(${OUT} ${PYTHON_OUT} PARENT_SCOPE) +endfunction() diff --git a/natten/cmake/get_gpu_lang.py b/natten/cmake/get_gpu_lang.py new file mode 100644 index 00000000..1eedff7e --- /dev/null +++ b/natten/cmake/get_gpu_lang.py @@ -0,0 +1,20 @@ +#!/usr/bin/env python3 + +import sys + +try: + import torch +except ImportError: + print("Torch is required for configuring a kernel build.", file=sys.stderr) + sys.exit(1) + +if torch.version.cuda is not None: + print("CUDA") +elif torch.version.hip is not None: + print("HIP") +elif torch.backends.mps.is_available(): + print("METAL") +elif hasattr(torch.version, "xpu") and torch.version.xpu is not None: + print("SYCL") +else: + print("CPU") diff --git a/natten/cmake/hipify.py b/natten/cmake/hipify.py new file mode 100644 index 00000000..a1539c02 --- /dev/null +++ b/natten/cmake/hipify.py @@ -0,0 +1,76 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 + +# From vLLM: https://github.com/vllm-project/vllm/blob/main/cmake/hipify.py + +# +# A command line tool for running pytorch's hipify preprocessor on CUDA +# source files. +# +# See https://github.com/ROCm/hipify_torch +# and /utils/hipify/hipify_python.py +# + +import argparse +import os +import shutil + +from torch.utils.hipify.hipify_python import hipify + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + + # Project directory where all the source + include files live. + parser.add_argument( + "-p", + "--project_dir", + help="The project directory.", + ) + + # Directory where hipified files are written. + parser.add_argument( + "-o", + "--output_dir", + help="The output directory.", + ) + + # Source files to convert. + parser.add_argument("sources", + help="Source files to hipify.", + nargs="*", + default=[]) + + args = parser.parse_args() + + # Limit include scope to project_dir only + includes = [os.path.join(args.project_dir, '*')] + + # Get absolute path for all source files. + extra_files = [os.path.abspath(s) for s in args.sources] + + # Copy sources from project directory to output directory. + # The directory might already exist to hold object files so we ignore that. + shutil.copytree(args.project_dir, args.output_dir, dirs_exist_ok=True) + + hipify_result = hipify(project_directory=args.project_dir, + output_directory=args.output_dir, + header_include_dirs=[], + includes=includes, + extra_files=extra_files, + show_detailed=True, + is_pytorch_extension=True, + hipify_extra_files_only=True) + + hipified_sources = [] + for source in args.sources: + s_abs = os.path.abspath(source) + hipified_s_abs = (hipify_result[s_abs].hipified_path if + (s_abs in hipify_result + and hipify_result[s_abs].hipified_path is not None) + else s_abs) + hipified_sources.append(hipified_s_abs) + + assert (len(hipified_sources) == len(args.sources)) + + # Print hipified source files. + print("\n".join(hipified_sources)) diff --git a/natten/cmake/kernel.cmake b/natten/cmake/kernel.cmake new file mode 100644 index 00000000..c41aaa65 --- /dev/null +++ b/natten/cmake/kernel.cmake @@ -0,0 +1,317 @@ +function(accumulate_gpu_archs OUT_ACC ACC EXTRA_ARCHS) + list(APPEND ACC ${EXTRA_ARCHS}) + list(REMOVE_DUPLICATES ACC) + list(SORT ACC) + set(${OUT_ACC} ${ACC} PARENT_SCOPE) +endfunction() + +function(cuda_kernel_component SRC_VAR) + set(oneValueArgs CUDA_MINVER NAME) + set(multiValueArgs SOURCES INCLUDES CUDA_CAPABILITIES CUDA_FLAGS CXX_FLAGS) + cmake_parse_arguments(KERNEL "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) + + if(NOT KERNEL_SOURCES) + message(FATAL_ERROR "cuda_kernel_component: SOURCES argument is required") + endif() + + # Bail out if this component is not supported by the CUDA version. + if(KERNEL_CUDA_MINVER) + if(CUDA_VERSION VERSION_LESS ${KERNEL_CUDA_MINVER}) + return() + endif() + endif() + + set(_KERNEL_SRC ${KERNEL_SOURCES}) + + if(KERNEL_INCLUDES) + # TODO: check if CLion support this: + # https://youtrack.jetbrains.com/issue/CPP-16510/CLion-does-not-handle-per-file-include-directories + set_source_files_properties( + ${_KERNEL_SRC} + PROPERTIES INCLUDE_DIRECTORIES "${KERNEL_INCLUDES}") + endif() + + # Determine CUDA architectures + if(KERNEL_CUDA_CAPABILITIES) + cuda_archs_loose_intersection(_KERNEL_ARCHS "${KERNEL_CUDA_CAPABILITIES}" "${CUDA_ARCHS}") + if(NOT _KERNEL_ARCHS) + message(FATAL_ERROR "CUDA kernel: ${KERNEL_NAME}, empty set of capabilities after intersection (kernel: ${KERNEL_CUDA_CAPABILITIES}, supported: ${CUDA_ARCHS})") + endif() + else() + set(_KERNEL_ARCHS "${CUDA_KERNEL_ARCHS}") + endif() + message(STATUS "CUDA kernel: ${KERNEL_NAME}, capabilities: ${_KERNEL_ARCHS}") + set_gencode_flags_for_srcs(SRCS "${_KERNEL_SRC}" CUDA_ARCHS "${_KERNEL_ARCHS}") + + accumulate_gpu_archs(_ALL_GPU_ARCHS "${ALL_GPU_ARCHS}" "${_KERNEL_ARCHS}") + set(ALL_GPU_ARCHS ${_ALL_GPU_ARCHS} PARENT_SCOPE) + + # Apply CUDA-specific compile flags + if(KERNEL_CUDA_FLAGS) + set(_CUDA_FLAGS "${KERNEL_CUDA_FLAGS}") + # -static-global-template-stub is not supported on CUDA < 12.8. Remove this + # once we don't support CUDA 12.6 anymore. + if(CUDA_VERSION VERSION_LESS 12.8) + string(REGEX REPLACE "-static-global-template-stub=(true|false)" "" _CUDA_FLAGS "${_CUDA_FLAGS}") + endif() + + foreach(_SRC ${_KERNEL_SRC}) + if(_SRC MATCHES ".*\\.cu$") + set_property( + SOURCE ${_SRC} + APPEND PROPERTY + COMPILE_OPTIONS "$<$:${_CUDA_FLAGS}>" + ) + endif() + endforeach() + endif() + + # Apply CXX-specific compile flags + if(KERNEL_CXX_FLAGS) + foreach(_SRC ${_KERNEL_SRC}) + set_property( + SOURCE ${_SRC} + APPEND PROPERTY + COMPILE_OPTIONS "$<$:${KERNEL_CXX_FLAGS}>" + ) + endforeach() + endif() + + set(_TMP_SRC ${${SRC_VAR}}) + list(APPEND _TMP_SRC ${_KERNEL_SRC}) + set(${SRC_VAR} ${_TMP_SRC} PARENT_SCOPE) +endfunction() + +function(hip_kernel_component SRC_VAR) + set(options SUPPORTS_HIPIFY) + set(oneValueArgs CUDA_MINVER NAME) + set(multiValueArgs SOURCES INCLUDES CXX_FLAGS HIP_FLAGS ROCM_ARCHS) + cmake_parse_arguments(KERNEL "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) + + if(NOT KERNEL_SOURCES) + message(FATAL_ERROR "hip_kernel_component: SOURCES argument is required") + endif() + + set(_KERNEL_SRC ${KERNEL_SOURCES}) + + if(KERNEL_INCLUDES) + # TODO: check if CLion support this: + # https://youtrack.jetbrains.com/issue/CPP-16510/CLion-does-not-handle-per-file-include-directories + set_source_files_properties( + ${_KERNEL_SRC} + PROPERTIES INCLUDE_DIRECTORIES "${KERNEL_INCLUDES}") + endif() + + # Apply HIP-specific compile flags + if(KERNEL_HIP_FLAGS) + foreach(_SRC ${_KERNEL_SRC}) + if(_SRC MATCHES ".*\\.(cu|hip)$") + set_property( + SOURCE ${_SRC} + APPEND PROPERTY + COMPILE_OPTIONS "$<$:${KERNEL_HIP_FLAGS}>" + ) + endif() + endforeach() + endif() + + # Determine ROCm architectures + if(KERNEL_ROCM_ARCHS) + hip_archs_loose_intersection(_KERNEL_ARCHS "${KERNEL_ROCM_ARCHS}" "${ROCM_ARCHS}") + if(NOT _KERNEL_ARCHS) + message(FATAL_ERROR "ROCm kernel: ${KERNEL_NAME}, empty set of architectures after intersection (kernel: ${KERNEL_ROCM_ARCHS}, supported: ${ROCM_ARCHS})") + endif() + else() + set(_KERNEL_ARCHS "${ROCM_ARCHS}") + endif() + message(STATUS "ROCm kernel: ${KERNEL_NAME}, archs: ${_KERNEL_ARCHS}") + + accumulate_gpu_archs(_ALL_GPU_ARCHS "${ALL_GPU_ARCHS}" "${_KERNEL_ARCHS}") + set(ALL_GPU_ARCHS ${_ALL_GPU_ARCHS} PARENT_SCOPE) + + foreach(_SRC ${_KERNEL_SRC}) + if(_SRC MATCHES ".*\\.(cu|hip)$") + foreach(_ARCH ${_KERNEL_ARCHS}) + set_property( + SOURCE ${_SRC} + APPEND PROPERTY + COMPILE_OPTIONS "$<$:--offload-arch=${_ARCH}>" + ) + endforeach() + endif() + endforeach() + + set(_TMP_SRC ${${SRC_VAR}}) + list(APPEND _TMP_SRC ${_KERNEL_SRC}) + set(${SRC_VAR} ${_TMP_SRC} PARENT_SCOPE) +endfunction() + + +function(xpu_kernel_component SRC_VAR) + set(options) + set(oneValueArgs) + set(multiValueArgs SOURCES INCLUDES CXX_FLAGS SYCL_FLAGS) + cmake_parse_arguments(KERNEL "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) + + if(NOT KERNEL_SOURCES) + message(FATAL_ERROR "xpu_kernel_component: SOURCES argument is required") + endif() + + set(_KERNEL_SRC ${KERNEL_SOURCES}) + + # Handle per-file include directories if specified + if(KERNEL_INCLUDES) + # TODO: check if CLion support this: + # https://youtrack.jetbrains.com/issue/CPP-16510/CLion-does-not-handle-per-file-include-directories + set_source_files_properties( + ${_KERNEL_SRC} + PROPERTIES INCLUDE_DIRECTORIES "${KERNEL_INCLUDES}") + endif() + + # Apply CXX-specific compile flags + if(KERNEL_CXX_FLAGS) + foreach(_SRC ${_KERNEL_SRC}) + set_property( + SOURCE ${_SRC} + APPEND PROPERTY + COMPILE_OPTIONS "$<$:${KERNEL_CXX_FLAGS}>" + ) + endforeach() + endif() + + # Add SYCL-specific compilation flags for XPU sources + if(KERNEL_SYCL_FLAGS) + # Use kernel-specific SYCL flags plus SYCL flags from the parent scope. + set(_SYCL_FLAGS ${sycl_flags}) + list(APPEND _SYCL_FLAGS ${KERNEL_SYCL_FLAGS}) + foreach(_SRC ${_KERNEL_SRC}) + if(_SRC MATCHES ".*\\.(cpp|cxx|cc)$") + set_property( + SOURCE ${_SRC} + APPEND PROPERTY + COMPILE_OPTIONS "$<$:${_SYCL_FLAGS}>" + ) + endif() + endforeach() + else() + # Use default SYCL flags (from parent scope variable sycl_flags) + foreach(_SRC ${_KERNEL_SRC}) + if(_SRC MATCHES ".*\\.(cpp|cxx|cc)$") + set_property( + SOURCE ${_SRC} + APPEND PROPERTY + COMPILE_OPTIONS "$<$:${sycl_flags}>" + ) + endif() + endforeach() + endif() + + # Append to parent scope SRC variable + set(_TMP_SRC ${${SRC_VAR}}) + list(APPEND _TMP_SRC ${_KERNEL_SRC}) + set(${SRC_VAR} ${_TMP_SRC} PARENT_SCOPE) +endfunction() + +function(cpu_kernel_component SRC_VAR) + set(options) + set(oneValueArgs) + set(multiValueArgs SOURCES INCLUDES CXX_FLAGS) + cmake_parse_arguments(KERNEL "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) + + if(NOT KERNEL_SOURCES) + message(FATAL_ERROR "cpu_kernel_component: SOURCES argument is required") + endif() + + set(_KERNEL_SRC ${KERNEL_SOURCES}) + + # Handle per-file include directories if specified + if(KERNEL_INCLUDES) + # TODO: check if CLion support this: + # https://youtrack.jetbrains.com/issue/CPP-16510/CLion-does-not-handle-per-file-include-directories + set_source_files_properties( + ${_KERNEL_SRC} + PROPERTIES INCLUDE_DIRECTORIES "${KERNEL_INCLUDES}") + endif() + + # Apply CXX-specific compile flags + if(KERNEL_CXX_FLAGS) + foreach(_SRC ${_KERNEL_SRC}) + set_property( + SOURCE ${_SRC} + APPEND PROPERTY + COMPILE_OPTIONS "$<$:${KERNEL_CXX_FLAGS}>" + ) + endforeach() + endif() + + # Append to parent scope SRC variable + set(_TMP_SRC ${${SRC_VAR}}) + list(APPEND _TMP_SRC ${_KERNEL_SRC}) + set(${SRC_VAR} ${_TMP_SRC} PARENT_SCOPE) +endfunction() + +function(metal_kernel_component SRC_VAR) + set(options) + set(oneValueArgs) + set(multiValueArgs SOURCES INCLUDES CXX_FLAGS) + cmake_parse_arguments(KERNEL "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) + + if(NOT KERNEL_SOURCES) + message(FATAL_ERROR "metal_kernel_component: SOURCES argument is required") + endif() + + set(_KERNEL_SRC ${KERNEL_SOURCES}) + + # Separate Metal shader files from other sources + set(_METAL_SRC) + set(_CPP_SRC) + + foreach(_SRC_FILE IN LISTS _KERNEL_SRC) + if(_SRC_FILE MATCHES "\\.(metal|h)$") + list(APPEND _METAL_SRC ${_SRC_FILE}) + else() + list(APPEND _CPP_SRC ${_SRC_FILE}) + endif() + endforeach() + + # Handle per-file include directories if specified (for C++ sources only) + if(KERNEL_INCLUDES AND _CPP_SRC) + # TODO: check if CLion support this: + # https://youtrack.jetbrains.com/issue/CPP-16510/CLion-does-not-handle-per-file-include-directories + set_source_files_properties( + ${_CPP_SRC} + PROPERTIES INCLUDE_DIRECTORIES "${KERNEL_INCLUDES}") + endif() + + # Apply CXX-specific compile flags + if(KERNEL_CXX_FLAGS AND _CPP_SRC) + foreach(_SRC ${_CPP_SRC}) + set_property( + SOURCE ${_SRC} + APPEND PROPERTY + COMPILE_OPTIONS "$<$:${KERNEL_CXX_FLAGS}>" + ) + endforeach() + endif() + + # Add C++ sources to main source list + if(_CPP_SRC) + set(_TMP_SRC ${${SRC_VAR}}) + list(APPEND _TMP_SRC ${_CPP_SRC}) + set(${SRC_VAR} ${_TMP_SRC} PARENT_SCOPE) + endif() + + # Keep track of Metal sources for later compilation + if(_METAL_SRC) + set(_TMP_METAL ${ALL_METAL_SOURCES}) + list(APPEND _TMP_METAL ${_METAL_SRC}) + set(ALL_METAL_SOURCES ${_TMP_METAL} PARENT_SCOPE) + endif() + + # Keep the includes directory for the Metal sources + if(KERNEL_INCLUDES AND _METAL_SRC) + set(_TMP_METAL_INCLUDES ${METAL_INCLUDE_DIRS}) + list(APPEND _TMP_METAL_INCLUDES ${KERNEL_INCLUDES}) + set(METAL_INCLUDE_DIRS ${_TMP_METAL_INCLUDES} PARENT_SCOPE) + endif() +endfunction() diff --git a/natten/cmake/metallib_to_header.py b/natten/cmake/metallib_to_header.py new file mode 100644 index 00000000..82bd252b --- /dev/null +++ b/natten/cmake/metallib_to_header.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python3 +import sys +import os + +def convert_metallib_to_header(metallib_path: str, header_path: str, target_name: str) -> None: + """Convert a metallib binary file to a C++ header with embedded data.""" + + # Read the metallib binary data + with open(metallib_path, 'rb') as f: + data: bytes = f.read() + + # Generate the header content + header_content: str = """// Auto-generated file containing embedded Metal library +#pragma once +#include +#include + +namespace """ + target_name + """_metal { + static const unsigned char metallib_data[] = { +""" + + # Convert binary data to C array format + bytes_per_line: int = 16 + for i in range(0, len(data), bytes_per_line): + chunk: bytes = data[i:i + bytes_per_line] + hex_values: str = ', '.join('0x{:02x}'.format(b) for b in chunk) + header_content += " " + hex_values + "," + if i + bytes_per_line < len(data): + header_content += "\n" + + header_content += """ + }; + static const size_t metallib_data_len = """ + str(len(data)) + """; + + // Convenience function to create Metal library from embedded data + inline id createLibrary(id device, NSError** error = nullptr) { + dispatch_data_t libraryData = dispatch_data_create( + metallib_data, + metallib_data_len, + dispatch_get_main_queue(), + ^{ /* No cleanup needed for static data */ }); + + NSError* localError = nil; + id library = [device newLibraryWithData:libraryData error:&localError]; + + if (error) { + *error = localError; + } + + return library; + } +} // namespace """ + target_name + """_metal +""" + + # Write the header file + dir_path: str = os.path.dirname(header_path) + if dir_path: + os.makedirs(dir_path, exist_ok=True) + with open(header_path, 'w') as f: + f.write(header_content) + + print("Generated {} ({} bytes)".format(header_path, len(data))) + +if __name__ == "__main__": + if len(sys.argv) != 4: + print("Usage: metallib_to_header.py ") + sys.exit(1) + + metallib_path: str = sys.argv[1] + header_path: str = sys.argv[2] + target_name: str = sys.argv[3] + + convert_metallib_to_header(metallib_path, header_path, target_name) \ No newline at end of file diff --git a/natten/cmake/utils.cmake b/natten/cmake/utils.cmake new file mode 100644 index 00000000..4c19c65c --- /dev/null +++ b/natten/cmake/utils.cmake @@ -0,0 +1,611 @@ +# Vendored from vLLM: +# +# https://github.com/vllm-project/vllm/blob/main/cmake/utils.cmake +# +# Attempt to find the python package that uses the same python executable as +# `EXECUTABLE` and is one of the `SUPPORTED_VERSIONS`. +# +macro (find_python_from_executable EXECUTABLE SUPPORTED_VERSIONS) + file(REAL_PATH ${EXECUTABLE} EXECUTABLE) + set(Python3_EXECUTABLE ${EXECUTABLE}) + find_package(Python3 COMPONENTS Interpreter Development.Module Development.SABIModule) + if (NOT Python3_FOUND) + message(FATAL_ERROR "Unable to find python matching: ${EXECUTABLE}.") + endif() + set(_VER "${Python3_VERSION_MAJOR}.${Python3_VERSION_MINOR}") + set(_SUPPORTED_VERSIONS_LIST ${SUPPORTED_VERSIONS} ${ARGN}) + if (NOT _VER IN_LIST _SUPPORTED_VERSIONS_LIST) + message(FATAL_ERROR + "Python version (${_VER}) is not one of the supported versions: " + "${_SUPPORTED_VERSIONS_LIST}.") + endif() + message(STATUS "Found python matching: ${EXECUTABLE}.") +endmacro() + +# +# Run `EXPR` in python. The standard output of python is stored in `OUT` and +# has trailing whitespace stripped. If an error is encountered when running +# python, a fatal message `ERR_MSG` is issued. +# +function (run_python OUT EXPR ERR_MSG) + if(Python3_EXECUTABLE) + set(_PYTHON_EXECUTABLE "${Python3_EXECUTABLE}") + elseif(Python_EXECUTABLE) + set(_PYTHON_EXECUTABLE "${Python_EXECUTABLE}") + else() + message(FATAL_ERROR "No Python executable found. Set Python3_EXECUTABLE or Python_EXECUTABLE.") + endif() + execute_process( + COMMAND + "${_PYTHON_EXECUTABLE}" "-c" "${EXPR}" + OUTPUT_VARIABLE PYTHON_OUT + RESULT_VARIABLE PYTHON_ERROR_CODE + ERROR_VARIABLE PYTHON_STDERR + OUTPUT_STRIP_TRAILING_WHITESPACE) + + if(NOT PYTHON_ERROR_CODE EQUAL 0) + message(FATAL_ERROR "${ERR_MSG}: ${PYTHON_STDERR}") + endif() + set(${OUT} ${PYTHON_OUT} PARENT_SCOPE) +endfunction() + +# +# Run `SCRIPT_PATH` in Python. The standard output of Python is stored in +# `OUT` and has trailing whitespace stripped. If the script exits with a +# non-zero code, a fatal message `ERR_MSG` is issued. +# +function (run_python_script OUT SCRIPT_PATH ERR_MSG) + if(Python3_EXECUTABLE) + set(_PYTHON_EXECUTABLE "${Python3_EXECUTABLE}") + elseif(Python_EXECUTABLE) + set(_PYTHON_EXECUTABLE "${Python_EXECUTABLE}") + else() + message(FATAL_ERROR "No Python executable found. Set Python3_EXECUTABLE or Python_EXECUTABLE.") + endif() + execute_process( + COMMAND + "${_PYTHON_EXECUTABLE}" "${SCRIPT_PATH}" + OUTPUT_VARIABLE PYTHON_OUT + RESULT_VARIABLE PYTHON_ERROR_CODE + ERROR_VARIABLE PYTHON_STDERR + OUTPUT_STRIP_TRAILING_WHITESPACE) + + if(NOT PYTHON_ERROR_CODE EQUAL 0) + message(FATAL_ERROR "${ERR_MSG}: ${PYTHON_STDERR}") + endif() + set(${OUT} ${PYTHON_OUT} PARENT_SCOPE) +endfunction() + +# +# Run `EXPR` in python. The standard output of python is stored in `OUT` and +# has trailing whitespace stripped. If an error is encountered when running +# python, `SUCCESS` is set to FALSE. If successful, `SUCCESS` is set to TRUE. +# +function (try_run_python OUT SUCCESS EXPR) + if(Python3_EXECUTABLE) + set(_PYTHON_EXECUTABLE "${Python3_EXECUTABLE}") + elseif(Python_EXECUTABLE) + set(_PYTHON_EXECUTABLE "${Python_EXECUTABLE}") + else() + message(FATAL_ERROR "No Python executable found. Set Python3_EXECUTABLE or Python_EXECUTABLE.") + endif() + execute_process( + COMMAND + "${_PYTHON_EXECUTABLE}" "-c" "${EXPR}" + OUTPUT_VARIABLE PYTHON_OUT + RESULT_VARIABLE PYTHON_ERROR_CODE + ERROR_QUIET + OUTPUT_STRIP_TRAILING_WHITESPACE) + + if(NOT PYTHON_ERROR_CODE EQUAL 0) + set(${SUCCESS} FALSE PARENT_SCOPE) + set(${OUT} "" PARENT_SCOPE) + else() + set(${SUCCESS} TRUE PARENT_SCOPE) + set(${OUT} ${PYTHON_OUT} PARENT_SCOPE) + endif() +endfunction() + +# Run `EXPR` in python after importing `PKG`. Use the result of this to extend +# `CMAKE_PREFIX_PATH` so the torch cmake configuration can be imported. +macro (append_cmake_prefix_path PKG EXPR) + run_python(_PREFIX_PATH + "import ${PKG}; print(${EXPR})" "Failed to locate ${PKG} path") + list(APPEND CMAKE_PREFIX_PATH ${_PREFIX_PATH}) +endmacro() + +# +# Add a target named `hipify${NAME}` that runs the hipify preprocessor on a set +# of CUDA source files. The names of the corresponding "hipified" sources are +# stored in `OUT_SRCS`. +# +function (hipify_sources_target OUT_SRCS NAME ORIG_SRCS) + # + # Split into C++ and non-C++ (i.e. CUDA) sources. + # + set(NODUP_SRCS ${ORIG_SRCS}) + list(REMOVE_DUPLICATES NODUP_SRCS) + set(SRCS ${NODUP_SRCS}) + set(CXX_SRCS ${NODUP_SRCS}) + list(FILTER SRCS INCLUDE REGEX "\.cu$") + list(FILTER CXX_SRCS EXCLUDE REGEX "\.cu$") + + # + # Generate ROCm/HIP source file names from CUDA file names. + # Since HIP files are generated code, they will appear in the build area + # `CMAKE_CURRENT_BINARY_DIR` directory rather than the original csrc dir. + # + set(HIP_SRCS) + foreach (SRC ${SRCS}) + get_source_file_property(include_dirs "${SRC}" INCLUDE_DIRECTORIES) + get_source_file_property(compile_options "${SRC}" COMPILE_OPTIONS) + string(REGEX REPLACE "\.cu$" "\.hip" SRC ${SRC}) + string(REGEX REPLACE "cuda" "hip" SRC ${SRC}) + + if(include_dirs) + # Copy over include directories from the original CUDA file. + set_source_files_properties( + ${SRC} + PROPERTIES INCLUDE_DIRECTORIES "${include_dirs}") + endif() + + if(compile_options) + set_source_files_properties( + ${SRC} + PROPERTIES COMPILE_OPTIONS "${compile_options}") + endif() + + list(APPEND HIP_SRCS "${CMAKE_CURRENT_BINARY_DIR}/${SRC}") + endforeach() + + add_custom_target( + hipify${NAME} + COMMAND "${Python3_EXECUTABLE}" ${CMAKE_SOURCE_DIR}/cmake/hipify.py -p ${CMAKE_SOURCE_DIR} -o ${CMAKE_CURRENT_BINARY_DIR} ${SRCS} + DEPENDS ${CMAKE_SOURCE_DIR}/cmake/hipify.py ${SRCS} + BYPRODUCTS ${HIP_SRCS} + COMMENT "Running hipify on ${NAME} extension source files.") + + # Swap out original extension sources with hipified sources. + list(APPEND HIP_SRCS ${CXX_SRCS}) + set(${OUT_SRCS} ${HIP_SRCS} PARENT_SCOPE) +endfunction() + +# +# Get additional GPU compiler flags from torch. +# +function (get_torch_gpu_compiler_flags OUT_GPU_FLAGS GPU_LANG) + if (${GPU_LANG} STREQUAL "CUDA") + # + # Get common NVCC flags from torch. + # + run_python(GPU_FLAGS + "from torch.utils.cpp_extension import COMMON_NVCC_FLAGS; print(';'.join(COMMON_NVCC_FLAGS))" + "Failed to determine torch nvcc compiler flags") + + if (CUDA_VERSION VERSION_GREATER_EQUAL 11.8) + list(APPEND GPU_FLAGS "-DENABLE_FP8") + list(REMOVE_ITEM GPU_FLAGS + "-D__CUDA_NO_HALF_OPERATORS__" + "-D__CUDA_NO_HALF_CONVERSIONS__" + "-D__CUDA_NO_BFLOAT16_CONVERSIONS__" + "-D__CUDA_NO_HALF2_OPERATORS__") + endif() + + elseif(${GPU_LANG} STREQUAL "HIP") + # + # Get common HIP/HIPCC flags from torch. + # + run_python(GPU_FLAGS + "import torch.utils.cpp_extension as t; print(';'.join(t.COMMON_HIP_FLAGS + t.COMMON_HIPCC_FLAGS))" + "Failed to determine torch nvcc compiler flags") + + list(APPEND GPU_FLAGS + "-DUSE_ROCM" + "-DENABLE_FP8" + "-U__HIP_NO_HALF_CONVERSIONS__" + "-U__HIP_NO_HALF_OPERATORS__" + "-fno-gpu-rdc") + + endif() + set(${OUT_GPU_FLAGS} ${GPU_FLAGS} PARENT_SCOPE) +endfunction() + +# Macro for converting a `gencode` version number to a cmake version number. +macro(string_to_ver OUT_VER IN_STR) + string(REGEX REPLACE "\([0-9]+\)\([0-9]\)" "\\1.\\2" ${OUT_VER} ${IN_STR}) +endmacro() + +# +# Clear all `-gencode` flags from `CMAKE_CUDA_FLAGS`. +# +# Example: +# CMAKE_CUDA_FLAGS="-Wall -gencode arch=compute_70,code=sm_70 -gencode arch=compute_75,code=sm_75" +# clear_gencode_flags() +# CMAKE_CUDA_FLAGS="-Wall" +# +macro(clear_gencode_flags) + # Remove all `-gencode` flags from `CMAKE_CUDA_FLAGS` since they will be modified + # and passed back via the `CUDA_ARCHITECTURES` property. + string(REGEX REPLACE "-gencode arch=[^ ]+ *" "" CMAKE_CUDA_FLAGS + ${CMAKE_CUDA_FLAGS}) +endmacro() + +# +# Extract unique CUDA architectures from a list of compute capabilities codes in +# the form `[]`, convert them to the form sort +# `.`, dedupes them and then sorts them in ascending order and +# stores them in `OUT_ARCHES`. +# +# Example: +# CUDA_ARCH_FLAGS="-gencode arch=compute_75,code=sm_75;...;-gencode arch=compute_90a,code=sm_90a" +# extract_unique_cuda_archs_ascending(OUT_ARCHES CUDA_ARCH_FLAGS) +# OUT_ARCHES="7.5;...;9.0" +function(extract_unique_cuda_archs_ascending OUT_ARCHES CUDA_ARCH_FLAGS) + set(_CUDA_ARCHES) + foreach(_ARCH ${CUDA_ARCH_FLAGS}) + string(REGEX MATCH "arch=compute_\([0-9]+a?\)" _COMPUTE ${_ARCH}) + if (_COMPUTE) + set(_COMPUTE ${CMAKE_MATCH_1}) + endif() + + string_to_ver(_COMPUTE_VER ${_COMPUTE}) + list(APPEND _CUDA_ARCHES ${_COMPUTE_VER}) + endforeach() + + list(REMOVE_DUPLICATES _CUDA_ARCHES) + list(SORT _CUDA_ARCHES COMPARE NATURAL ORDER ASCENDING) + set(${OUT_ARCHES} ${_CUDA_ARCHES} PARENT_SCOPE) +endfunction() + +# +# For a specific file set the `-gencode` flag in compile options conditionally +# for the CUDA language. +# +# Example: +# set_gencode_flag_for_srcs( +# SRCS "foo.cu" +# ARCH "compute_75" +# CODE "sm_75") +# adds: "-gencode arch=compute_75,code=sm_75" to the compile options for +# `foo.cu` (only for the CUDA language). +# +macro(set_gencode_flag_for_srcs) + set(options) + set(oneValueArgs ARCH CODE) + set(multiValueArgs SRCS) + cmake_parse_arguments(arg "${options}" "${oneValueArgs}" + "${multiValueArgs}" ${ARGN} ) + set(_FLAG -gencode arch=${arg_ARCH},code=${arg_CODE}) + set_property( + SOURCE ${arg_SRCS} + APPEND PROPERTY + COMPILE_OPTIONS "$<$:${_FLAG}>" + ) + + message(DEBUG "Setting gencode flag for ${arg_SRCS}: ${_FLAG}") +endmacro(set_gencode_flag_for_srcs) + +# +# For a list of source files set the `-gencode` flags in the files specific +# compile options (specifically for the CUDA language). +# +# arguments are: +# SRCS: list of source files +# CUDA_ARCHS: list of CUDA architectures in the form `.[letter]` +# BUILD_PTX_FOR_ARCH: if set to true, then the PTX code will be built +# for architecture `BUILD_PTX_FOR_ARCH` if there is a CUDA_ARCH in CUDA_ARCHS +# that is larger than BUILD_PTX_FOR_ARCH. +# +macro(set_gencode_flags_for_srcs) + set(options) + set(oneValueArgs BUILD_PTX_FOR_ARCH) + set(multiValueArgs SRCS CUDA_ARCHS) + cmake_parse_arguments(arg "${options}" "${oneValueArgs}" + "${multiValueArgs}" ${ARGN} ) + + foreach(_ARCH ${arg_CUDA_ARCHS}) + # handle +PTX suffix: generate both sm and ptx codes if requested + string(FIND "${_ARCH}" "+PTX" _HAS_PTX) + if(NOT _HAS_PTX EQUAL -1) + string(REPLACE "+PTX" "" _BASE_ARCH "${_ARCH}") + string(REPLACE "." "" _STRIPPED_ARCH "${_BASE_ARCH}") + set_gencode_flag_for_srcs( + SRCS ${arg_SRCS} + ARCH "compute_${_STRIPPED_ARCH}" + CODE "sm_${_STRIPPED_ARCH}") + set_gencode_flag_for_srcs( + SRCS ${arg_SRCS} + ARCH "compute_${_STRIPPED_ARCH}" + CODE "compute_${_STRIPPED_ARCH}") + else() + string(REPLACE "." "" _STRIPPED_ARCH "${_ARCH}") + set_gencode_flag_for_srcs( + SRCS ${arg_SRCS} + ARCH "compute_${_STRIPPED_ARCH}" + CODE "sm_${_STRIPPED_ARCH}") + endif() + endforeach() + + if (${arg_BUILD_PTX_FOR_ARCH}) + list(SORT arg_CUDA_ARCHS COMPARE NATURAL ORDER ASCENDING) + list(GET arg_CUDA_ARCHS -1 _HIGHEST_ARCH) + if (_HIGHEST_ARCH VERSION_GREATER_EQUAL ${arg_BUILD_PTX_FOR_ARCH}) + string(REPLACE "." "" _PTX_ARCH "${arg_BUILD_PTX_FOR_ARCH}") + set_gencode_flag_for_srcs( + SRCS ${arg_SRCS} + ARCH "compute_${_PTX_ARCH}" + CODE "compute_${_PTX_ARCH}") + endif() + endif() +endmacro() + +# +# For the given `SRC_CUDA_ARCHS` list of gencode versions in the form +# `.[letter]` compute the "loose intersection" with the +# `TGT_CUDA_ARCHS` list of gencodes. We also support the `+PTX` suffix in +# `SRC_CUDA_ARCHS` which indicates that the PTX code should be built when there +# is a CUDA_ARCH in `TGT_CUDA_ARCHS` that is equal to or larger than the +# architecture in `SRC_CUDA_ARCHS`. +# The loose intersection is defined as: +# { max{ x \in tgt | x <= y } | y \in src, { x \in tgt | x <= y } != {} } +# where `<=` is the version comparison operator. +# In other words, for each version in `TGT_CUDA_ARCHS` find the highest version +# in `SRC_CUDA_ARCHS` that is less or equal to the version in `TGT_CUDA_ARCHS`. +# We have special handling for x.0a, if x.0a is in `SRC_CUDA_ARCHS` and x.0 is +# in `TGT_CUDA_ARCHS` then we should remove x.0a from `SRC_CUDA_ARCHS` and add +# x.0a to the result (and remove x.0 from TGT_CUDA_ARCHS). +# The result is stored in `OUT_CUDA_ARCHS`. +# +# Example: +# SRC_CUDA_ARCHS="7.5;8.0;8.6;9.0;9.0a" +# TGT_CUDA_ARCHS="8.0;8.9;9.0" +# cuda_archs_loose_intersection(OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_ARCHS) +# OUT_CUDA_ARCHS="8.0;8.6;9.0;9.0a" +# +# Example With PTX: +# SRC_CUDA_ARCHS="8.0+PTX" +# TGT_CUDA_ARCHS="9.0" +# cuda_archs_loose_intersection(OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_ARCHS) +# OUT_CUDA_ARCHS="8.0+PTX" +# +function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_ARCHS) + set(_SRC_CUDA_ARCHS "${SRC_CUDA_ARCHS}") + set(_TGT_CUDA_ARCHS ${TGT_CUDA_ARCHS}) + + # handle +PTX suffix: separate base arch for matching, record PTX requests + set(_PTX_ARCHS) + foreach(_arch ${_SRC_CUDA_ARCHS}) + if(_arch MATCHES "\\+PTX$") + string(REPLACE "+PTX" "" _base "${_arch}") + list(APPEND _PTX_ARCHS "${_base}") + list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}") + list(APPEND _SRC_CUDA_ARCHS "${_base}") + endif() + endforeach() + list(REMOVE_DUPLICATES _PTX_ARCHS) + list(REMOVE_DUPLICATES _SRC_CUDA_ARCHS) + + # If x.0a or x.0f is in SRC_CUDA_ARCHS and x.0 is in CUDA_ARCHS then we should + # remove x.0a or x.0f from SRC_CUDA_ARCHS and add x.0a or x.0f to _CUDA_ARCHS + set(_CUDA_ARCHS) + foreach(_arch ${_SRC_CUDA_ARCHS}) + if(_arch MATCHES "[af]$") + list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}") + string(REGEX REPLACE "[af]$" "" _base "${_arch}") + if ("${_base}" IN_LIST TGT_CUDA_ARCHS) + list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}") + list(APPEND _CUDA_ARCHS "${_arch}") + endif() + endif() + endforeach() + + list(SORT _SRC_CUDA_ARCHS COMPARE NATURAL ORDER ASCENDING) + + # for each ARCH in TGT_CUDA_ARCHS find the highest arch in SRC_CUDA_ARCHS that + # is less or equal to ARCH (but has the same major version since SASS binary + # compatibility is only forward compatible within the same major version). + foreach(_ARCH ${_TGT_CUDA_ARCHS}) + set(_TMP_ARCH) + # Extract the major version of the target arch + string(REGEX REPLACE "^([0-9]+)\\..*$" "\\1" TGT_ARCH_MAJOR "${_ARCH}") + foreach(_SRC_ARCH ${_SRC_CUDA_ARCHS}) + # Extract the major version of the source arch + string(REGEX REPLACE "^([0-9]+)\\..*$" "\\1" SRC_ARCH_MAJOR "${_SRC_ARCH}") + # Check version-less-or-equal, and allow PTX arches to match across majors + if (_SRC_ARCH VERSION_LESS_EQUAL _ARCH) + if (_SRC_ARCH IN_LIST _PTX_ARCHS OR SRC_ARCH_MAJOR STREQUAL TGT_ARCH_MAJOR) + set(_TMP_ARCH "${_SRC_ARCH}") + endif() + else() + # If we hit a version greater than the target, we can break + break() + endif() + endforeach() + + # If we found a matching _TMP_ARCH, append it to _CUDA_ARCHS + if (_TMP_ARCH) + list(APPEND _CUDA_ARCHS "${_TMP_ARCH}") + endif() + endforeach() + + list(REMOVE_DUPLICATES _CUDA_ARCHS) + + # reapply +PTX suffix to architectures that requested PTX + set(_FINAL_ARCHS) + foreach(_arch ${_CUDA_ARCHS}) + if(_arch IN_LIST _PTX_ARCHS) + list(APPEND _FINAL_ARCHS "${_arch}+PTX") + else() + list(APPEND _FINAL_ARCHS "${_arch}") + endif() + endforeach() + set(_CUDA_ARCHS ${_FINAL_ARCHS}) + + list(SORT _CUDA_ARCHS COMPARE NATURAL ORDER ASCENDING) + + set(${OUT_CUDA_ARCHS} ${_CUDA_ARCHS} PARENT_SCOPE) +endfunction() + +# +# For the given `SRC_ROCM_ARCHS` list of architecture versions in the form +# `` compute the "loose intersection" with the `TGT_ROCM_ARCHS` list. +# The loose intersection is defined as: +# { max{ x \in tgt | x <= y } | y \in src, { x \in tgt | x <= y } != {} } +# where `<=` is the version comparison operator. +# In other words, for each version in `TGT_ROCM_ARCHS` find the highest version +# in `SRC_ROCM_ARCHS` that is less or equal to the version in `TGT_ROCM_ARCHS`. +# The result is stored in `OUT_ROCM_ARCHS`. +# +# Example: +# SRC_ROCM_ARCHS="gfx900;gfx906;gfx908;gfx90a" +# TGT_ROCM_ARCHS="gfx906;gfx908;gfx1030" +# hip_archs_loose_intersection(OUT_ROCM_ARCHS SRC_ROCM_ARCHS TGT_ROCM_ARCHS) +# OUT_ROCM_ARCHS="gfx906;gfx908" +# +function(hip_archs_loose_intersection OUT_ROCM_ARCHS SRC_ROCM_ARCHS TGT_ROCM_ARCHS) + list(REMOVE_DUPLICATES SRC_ROCM_ARCHS) + + # ROCm architectures are typically in format gfxNNN or gfxNNNx where N is a digit + # and x is a letter. We can sort them by string comparison which works for this format. + list(SORT SRC_ROCM_ARCHS COMPARE STRING ORDER ASCENDING) + + set(_ROCM_ARCHS) + + # Find the intersection of supported architectures + foreach(_SRC_ARCH ${SRC_ROCM_ARCHS}) + if(_SRC_ARCH IN_LIST TGT_ROCM_ARCHS) + list(APPEND _ROCM_ARCHS ${_SRC_ARCH}) + endif() + endforeach() + + list(REMOVE_DUPLICATES _ROCM_ARCHS) + set(${OUT_ROCM_ARCHS} ${_ROCM_ARCHS} PARENT_SCOPE) +endfunction() + +function(cuda_remove_ptx_suffixes OUT_CUDA_ARCHS CUDA_ARCHS) + set(_CUDA_ARCHS "${CUDA_ARCHS}") + + # handle +PTX suffix: separate base arch for matching, record PTX requests + foreach(_arch ${CUDA_ARCHS}) + if(_arch MATCHES "\\+PTX$") + string(REPLACE "+PTX" "" _base "${_arch}") + list(REMOVE_ITEM _CUDA_ARCHS "${_arch}") + list(APPEND _CUDA_ARCHS "${_base}") + endif() + endforeach() + + list(REMOVE_DUPLICATES _CUDA_ARCHS) + list(SORT _CUDA_ARCHS COMPARE NATURAL ORDER ASCENDING) + + set(${OUT_CUDA_ARCHS} ${_CUDA_ARCHS} PARENT_SCOPE) +endfunction() + +# +# Define a target named `GPU_MOD_NAME` for a single extension. The +# arguments are: +# +# DESTINATION - Module destination directory. +# LANGUAGE - The GPU language for this module, e.g CUDA, HIP, +# etc. +# SOURCES - List of source files relative to CMakeLists.txt +# directory. +# +# Optional arguments: +# +# ARCHITECTURES - A list of target GPU architectures in cmake +# format. +# Refer `CMAKE_CUDA_ARCHITECTURES` documentation +# and `CMAKE_HIP_ARCHITECTURES` for more info. +# ARCHITECTURES will use cmake's defaults if +# not provided. +# COMPILE_FLAGS - Extra compiler flags passed to NVCC/hip. +# INCLUDE_DIRECTORIES - Extra include directories. +# LIBRARIES - Extra link libraries. +# WITH_SOABI - Generate library with python SOABI suffix name. +# USE_SABI - Use python stable api +# +# Note: optimization level/debug info is set via cmake build type. +# +function (define_gpu_extension_target GPU_MOD_NAME) + cmake_parse_arguments(PARSE_ARGV 1 + GPU + "WITH_SOABI" + "DESTINATION;LANGUAGE;USE_SABI" + "SOURCES;COMPILE_FLAGS;INCLUDE_DIRECTORIES;LIBRARIES") + + # Add hipify preprocessing step when building with HIP/ROCm. + if (GPU_LANGUAGE STREQUAL "HIP") + hipify_sources_target(GPU_SOURCES ${GPU_MOD_NAME} "${GPU_SOURCES}") + endif() + + if (GPU_WITH_SOABI) + set(GPU_WITH_SOABI WITH_SOABI) + else() + set(GPU_WITH_SOABI) + endif() + + if (GPU_USE_SABI) + Python3_add_library(${GPU_MOD_NAME} MODULE USE_SABI ${GPU_USE_SABI} ${GPU_WITH_SOABI} "${GPU_SOURCES}") + else() + Python3_add_library(${GPU_MOD_NAME} MODULE ${GPU_WITH_SOABI} "${GPU_SOURCES}") + endif() + + if (GPU_LANGUAGE STREQUAL "HIP") + # Make this target dependent on the hipify preprocessor step. + add_dependencies(${GPU_MOD_NAME} hipify${GPU_MOD_NAME}) + + # Clear target architectures, we are passing arch flags per source file. + set_property(TARGET ${GPU_MOD_NAME} PROPERTY HIP_ARCHITECTURES off) + endif() + + if (TORCH_VERSION VERSION_LESS 2.12.0) + set_property(TARGET ${GPU_MOD_NAME} PROPERTY CXX_STANDARD 17) + else() + set_property(TARGET ${GPU_MOD_NAME} PROPERTY CXX_STANDARD 20) + endif() + + target_compile_options(${GPU_MOD_NAME} PRIVATE + $<$:${GPU_COMPILE_FLAGS}>) + + target_compile_definitions(${GPU_MOD_NAME} PRIVATE + "-DTORCH_EXTENSION_NAME=${GPU_MOD_NAME}") + + target_include_directories(${GPU_MOD_NAME} PRIVATE csrc + ${GPU_INCLUDE_DIRECTORIES}) + + target_link_libraries(${GPU_MOD_NAME} PRIVATE torch ${GPU_LIBRARIES}) + + # Don't use `TORCH_LIBRARIES` for CUDA since it pulls in a bunch of + # dependencies that are not necessary and may not be installed. + if (GPU_LANGUAGE STREQUAL "CUDA") + target_link_libraries(${GPU_MOD_NAME} PRIVATE CUDA::cudart) + else() + target_link_libraries(${GPU_MOD_NAME} PRIVATE ${TORCH_LIBRARIES}) + endif() + + install(TARGETS ${GPU_MOD_NAME} LIBRARY DESTINATION ${GPU_DESTINATION} COMPONENT ${GPU_MOD_NAME}) +endfunction() + +# Map a GPU language to its backend name. +# +# Arguments: +# OUT_BACKEND - Output variable name for the backend string +# GPU_LANG - The GPU language (CPU, CUDA, HIP, METAL, SYCL) +# +function(gpu_lang_to_backend OUT_BACKEND GPU_LANG) + if (${GPU_LANG} STREQUAL "CPU") + set(_BACKEND "cpu") + elseif (${GPU_LANG} STREQUAL "CUDA") + set(_BACKEND "cuda") + elseif (${GPU_LANG} STREQUAL "HIP") + set(_BACKEND "rocm") + elseif (${GPU_LANG} STREQUAL "METAL") + set(_BACKEND "metal") + elseif (${GPU_LANG} STREQUAL "SYCL") + set(_BACKEND "xpu") + else() + message(FATAL_ERROR "Unsupported GPU_LANG: ${GPU_LANG}") + endif() + + set(${OUT_BACKEND} "${_BACKEND}" PARENT_SCOPE) +endfunction() diff --git a/natten/compat.py b/natten/compat.py new file mode 100644 index 00000000..a9b2672c --- /dev/null +++ b/natten/compat.py @@ -0,0 +1,26 @@ +import ctypes +import importlib.util +import sys +from pathlib import Path +from types import ModuleType + + +def _import_from_path(file_path: Path) -> ModuleType: + # We cannot use the module name as-is, after adding it to `sys.modules`, + # it would also be used for other imports. So, we make a module name that + # depends on the path for it to be unique using the hex-encoded hash of + # the path. + path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) + module_name = path_hash + spec = importlib.util.spec_from_file_location(module_name, file_path) + if spec is None: + raise ImportError(f"Cannot load spec for {module_name} from {file_path}") + module = importlib.util.module_from_spec(spec) + if module is None: + raise ImportError(f"Cannot load module {module_name} from spec") + sys.modules[module_name] = module + spec.loader.exec_module(module) # type: ignore + return module + + +globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/natten/metadata-cann.json b/natten/metadata-cann.json new file mode 100644 index 00000000..3123bb7b --- /dev/null +++ b/natten/metadata-cann.json @@ -0,0 +1,22 @@ +{ + "name": "natten", + "id": "_natten_cann_28fa1dd", + "version": 1, + "license": "MIT", + "upstream": "https://github.com/SHI-Labs/NATTEN", + "python-depends": [], + "backend": { + "type": "cann" + }, + "provenance": { + "kernel-builder": { + "version": "0.17.0-dev0", + "sha": "81580bb92577f2f7228661ca2a221fc052375709", + "dirty": true + }, + "kernel": { + "sha": "39160cf398fdf251faf03ede7fa8d56dfa29701e", + "dirty": false + } + } +} \ No newline at end of file diff --git a/natten/metadata-cpu.json b/natten/metadata-cpu.json new file mode 100644 index 00000000..68482e01 --- /dev/null +++ b/natten/metadata-cpu.json @@ -0,0 +1,22 @@ +{ + "name": "natten", + "id": "_natten_cpu_28fa1dd", + "version": 1, + "license": "MIT", + "upstream": "https://github.com/SHI-Labs/NATTEN", + "python-depends": [], + "backend": { + "type": "cpu" + }, + "provenance": { + "kernel-builder": { + "version": "0.17.0-dev0", + "sha": "81580bb92577f2f7228661ca2a221fc052375709", + "dirty": true + }, + "kernel": { + "sha": "39160cf398fdf251faf03ede7fa8d56dfa29701e", + "dirty": false + } + } +} \ No newline at end of file diff --git a/natten/metadata-cuda.json b/natten/metadata-cuda.json new file mode 100644 index 00000000..5771d35d --- /dev/null +++ b/natten/metadata-cuda.json @@ -0,0 +1,22 @@ +{ + "name": "natten", + "id": "_natten_cuda_28fa1dd", + "version": 1, + "license": "MIT", + "upstream": "https://github.com/SHI-Labs/NATTEN", + "python-depends": [], + "backend": { + "type": "cuda" + }, + "provenance": { + "kernel-builder": { + "version": "0.17.0-dev0", + "sha": "81580bb92577f2f7228661ca2a221fc052375709", + "dirty": true + }, + "kernel": { + "sha": "39160cf398fdf251faf03ede7fa8d56dfa29701e", + "dirty": false + } + } +} \ No newline at end of file diff --git a/natten/metadata-metal.json b/natten/metadata-metal.json new file mode 100644 index 00000000..7fd48375 --- /dev/null +++ b/natten/metadata-metal.json @@ -0,0 +1,22 @@ +{ + "name": "natten", + "id": "_natten_metal_28fa1dd", + "version": 1, + "license": "MIT", + "upstream": "https://github.com/SHI-Labs/NATTEN", + "python-depends": [], + "backend": { + "type": "metal" + }, + "provenance": { + "kernel-builder": { + "version": "0.17.0-dev0", + "sha": "81580bb92577f2f7228661ca2a221fc052375709", + "dirty": true + }, + "kernel": { + "sha": "39160cf398fdf251faf03ede7fa8d56dfa29701e", + "dirty": false + } + } +} \ No newline at end of file diff --git a/natten/metadata-neuron.json b/natten/metadata-neuron.json new file mode 100644 index 00000000..7bed17df --- /dev/null +++ b/natten/metadata-neuron.json @@ -0,0 +1,22 @@ +{ + "name": "natten", + "id": "_natten_neuron_28fa1dd", + "version": 1, + "license": "MIT", + "upstream": "https://github.com/SHI-Labs/NATTEN", + "python-depends": [], + "backend": { + "type": "neuron" + }, + "provenance": { + "kernel-builder": { + "version": "0.17.0-dev0", + "sha": "81580bb92577f2f7228661ca2a221fc052375709", + "dirty": true + }, + "kernel": { + "sha": "39160cf398fdf251faf03ede7fa8d56dfa29701e", + "dirty": false + } + } +} \ No newline at end of file diff --git a/natten/metadata-rocm.json b/natten/metadata-rocm.json new file mode 100644 index 00000000..990b7bf9 --- /dev/null +++ b/natten/metadata-rocm.json @@ -0,0 +1,22 @@ +{ + "name": "natten", + "id": "_natten_rocm_28fa1dd", + "version": 1, + "license": "MIT", + "upstream": "https://github.com/SHI-Labs/NATTEN", + "python-depends": [], + "backend": { + "type": "rocm" + }, + "provenance": { + "kernel-builder": { + "version": "0.17.0-dev0", + "sha": "81580bb92577f2f7228661ca2a221fc052375709", + "dirty": true + }, + "kernel": { + "sha": "39160cf398fdf251faf03ede7fa8d56dfa29701e", + "dirty": false + } + } +} \ No newline at end of file diff --git a/natten/metadata-xpu.json b/natten/metadata-xpu.json new file mode 100644 index 00000000..8409363a --- /dev/null +++ b/natten/metadata-xpu.json @@ -0,0 +1,22 @@ +{ + "name": "natten", + "id": "_natten_xpu_28fa1dd", + "version": 1, + "license": "MIT", + "upstream": "https://github.com/SHI-Labs/NATTEN", + "python-depends": [], + "backend": { + "type": "xpu" + }, + "provenance": { + "kernel-builder": { + "version": "0.17.0-dev0", + "sha": "81580bb92577f2f7228661ca2a221fc052375709", + "dirty": true + }, + "kernel": { + "sha": "39160cf398fdf251faf03ede7fa8d56dfa29701e", + "dirty": false + } + } +} \ No newline at end of file diff --git a/natten/pyproject.toml b/natten/pyproject.toml new file mode 100644 index 00000000..3743f4d9 --- /dev/null +++ b/natten/pyproject.toml @@ -0,0 +1,23 @@ +[project] +name = "natten" +version = "0.1.0" +requires-python = ">=3.9" + +[build-system] +requires = [ + "cmake>=3.26", + "ninja", + "packaging", + "setuptools>=61", + "torch", + "wheel", + +] +build-backend = "setuptools.build_meta" + +[project.optional-dependencies] + +[tool.pytest.ini_options] +markers = [ + "kernels_ci: mark a test as a kernel CI test" +] \ No newline at end of file diff --git a/natten/setup.py b/natten/setup.py new file mode 100644 index 00000000..92f797d7 --- /dev/null +++ b/natten/setup.py @@ -0,0 +1,200 @@ +import logging +import os +from shutil import which, move +import subprocess +import sys +from pathlib import Path + +from setuptools import Extension, find_packages, setup +from setuptools.command.build import build +from setuptools.command.build_ext import build_ext + +logger = logging.getLogger(__name__) + + +def get_backend() -> str: + """Detect the backend by inspecting torch.""" + import torch + + if torch.version.cuda is not None: + return "cuda" + elif torch.version.hip is not None: + return "rocm" + elif torch.backends.mps.is_available(): + return "metal" + elif hasattr(torch.version, "xpu") and torch.version.xpu is not None: + return "xpu" + else: + return "cpu" + + +def is_sccache_available() -> bool: + return which("sccache") is not None + + +def is_ccache_available() -> bool: + return which("ccache") is not None + + +def is_ninja_available() -> bool: + return which("ninja") is not None + + +def _make_cmake_args(cfg: str) -> tuple[list[str], list[str]]: + """Build CMake and build arguments from the current environment.""" + cmake_generator = os.environ.get("CMAKE_GENERATOR", "") + + cmake_args = [ + f"-DPython3_EXECUTABLE={sys.executable}", + f"-DCMAKE_BUILD_TYPE={cfg}", # not used on MSVC, but no harm + ] + build_args: list[str] = [] + + if "CMAKE_ARGS" in os.environ: + cmake_args += [item for item in os.environ["CMAKE_ARGS"].split(" ") if item] + + if not cmake_generator or cmake_generator == "Ninja": + try: + import ninja + + ninja_executable_path = Path(ninja.BIN_DIR) / "ninja" + cmake_args += [ + "-GNinja", + f"-DCMAKE_MAKE_PROGRAM:FILEPATH={ninja_executable_path}", + ] + except ImportError: + pass + + if is_sccache_available(): + cmake_args += [ + "-DCMAKE_C_COMPILER_LAUNCHER=sccache", + "-DCMAKE_CXX_COMPILER_LAUNCHER=sccache", + "-DCMAKE_CUDA_COMPILER_LAUNCHER=sccache", + "-DCMAKE_HIP_COMPILER_LAUNCHER=sccache", + "-DCMAKE_OBJC_COMPILER_LAUNCHER=sccache", + "-DCMAKE_OBJCXX_COMPILER_LAUNCHER=sccache", + ] + elif is_ccache_available(): + cmake_args += [ + "-DCMAKE_C_COMPILER_LAUNCHER=ccache", + "-DCMAKE_CXX_COMPILER_LAUNCHER=ccache", + "-DCMAKE_CUDA_COMPILER_LAUNCHER=ccache", + "-DCMAKE_HIP_COMPILER_LAUNCHER=ccache", + "-DCMAKE_OBJC_COMPILER_LAUNCHER=ccache", + "-DCMAKE_OBJCXX_COMPILER_LAUNCHER=ccache", + ] + + num_jobs = os.getenv("MAX_JOBS", None) + if num_jobs is not None: + num_jobs = int(num_jobs) + logger.info("Using MAX_JOBS=%d as the number of jobs.", num_jobs) + else: + try: + # os.sched_getaffinity() isn't universally available, so fall + # back to os.cpu_count() if we get an error here. + num_jobs = len(os.sched_getaffinity(0)) + except AttributeError: + num_jobs = os.cpu_count() + + nvcc_threads = os.getenv("NVCC_THREADS", None) + if nvcc_threads is not None: + nvcc_threads = int(nvcc_threads) + logger.info( + "Using NVCC_THREADS=%d as the number of nvcc threads.", nvcc_threads + ) + num_jobs = max(1, num_jobs // nvcc_threads) + cmake_args += ["-DNVCC_THREADS={}".format(nvcc_threads)] + + build_args += [f"-j{num_jobs}"] + if sys.platform == "win32": + build_args += ["--config", cfg] + + return cmake_args, build_args + + +class CMakeExtension(Extension): + def __init__(self, name: str, sourcedir: str = "") -> None: + super().__init__(name, sources=[], py_limited_api=True) + self.sourcedir = os.fspath(Path(sourcedir).resolve()) + + +class CMakeBuild(build_ext): + def build_extension(self, ext: CMakeExtension) -> None: + ext_fullpath = Path.cwd() / self.get_ext_fullpath(ext.name) + extdir = ext_fullpath.parent.resolve() + + debug = int(os.environ.get("DEBUG", 0)) if self.debug is None else self.debug + cfg = "Debug" if debug else "Release" + + cmake_args, build_args = _make_cmake_args(cfg) + cmake_args = [f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={extdir}{os.sep}"] + cmake_args + + build_temp = Path(self.build_temp) / ext.name + if not build_temp.exists(): + build_temp.mkdir(parents=True) + + subprocess.run( + ["cmake", "-S", ext.sourcedir, "-B", str(build_temp), *cmake_args], + cwd=build_temp, + check=True, + ) + subprocess.run( + ["cmake", "--build", str(build_temp), *build_args], cwd=build_temp, check=True + ) + + if sys.platform == "win32": + # Move the dylib one folder up for discovery. + for filename in os.listdir(extdir / cfg): + move(extdir / cfg / filename, extdir / filename) + + +class BuildKernel(build): + """Custom command to build and locally install the kernel.""" + + description = "Build the kernel and install via the local_install CMake target" + user_options = [] + + def initialize_options(self) -> None: + super().initialize_options() + + def finalize_options(self) -> None: + super().finalize_options() + + def run(self) -> None: + project_root = Path(__file__).parent + + debug = int(os.environ.get("DEBUG", 0)) + cfg = "Debug" if debug else "Release" + + cmake_args, build_args = _make_cmake_args(cfg) + + build_temp = project_root / "_cmake_build" + build_temp.mkdir(parents=True, exist_ok=True) + + subprocess.run( + ["cmake", "-S", str(project_root), "-B", str(build_temp), *cmake_args], + cwd=project_root, + check=True, + ) + subprocess.run( + ["cmake", "--build", str(build_temp), "--target", "local_install", *build_args], + cwd=project_root, + check=True, + ) + + +backend = get_backend() +ops_name = f"_natten_{backend}_28fa1dd" + +setup( + name="natten", + # The version is just a stub, it's not used by the final build artefact. + version="0.1.0", + ext_modules=[CMakeExtension(f"natten.{ops_name}")], + cmdclass={"build_ext": CMakeBuild, "build_kernel": BuildKernel}, + packages=find_packages(where="torch-ext", include=["natten*"]), + package_dir={"": "torch-ext"}, + zip_safe=False, + install_requires=["torch"], + python_requires=">=3.9", +) \ No newline at end of file diff --git a/natten/torch-ext/natten/__init__.py b/natten/torch-ext/natten/__init__.py index 37ddeac9..9e0a40d5 100644 --- a/natten/torch-ext/natten/__init__.py +++ b/natten/torch-ext/natten/__init__.py @@ -63,6 +63,18 @@ ) from .version import __version__ +# kernel-builder port: the package contents are installed flat into the build +# variant directory, so a module literally named `types` would shadow the +# standard library `types` module whenever that directory is on PYTHONPATH +# (e.g. kernel-builder test shells and CI runners), breaking interpreter +# startup. The module therefore lives in `_types`; alias it here so +# `natten.types` keeps working like upstream. +import sys as _sys + +from . import _types as types + +_sys.modules[__name__ + ".types"] = types + __all__ = [ "__version__", "NeighborhoodAttention1D", diff --git a/natten/torch-ext/natten/types.py b/natten/torch-ext/natten/_types.py similarity index 100% rename from natten/torch-ext/natten/types.py rename to natten/torch-ext/natten/_types.py diff --git a/natten/torch-ext/natten/backends/blackwell_fmha.py b/natten/torch-ext/natten/backends/blackwell_fmha.py index c2816300..9b36e79d 100644 --- a/natten/torch-ext/natten/backends/blackwell_fmha.py +++ b/natten/torch-ext/natten/backends/blackwell_fmha.py @@ -37,7 +37,7 @@ check_cutlass_blackwell_fmha_backward_config, check_cutlass_blackwell_fmha_forward_config, ) -from ..types import ( +from .._types import ( CutlassBlackwellFmhaBackwardConfigType, CutlassBlackwellFmhaForwardConfigType, NoneType, diff --git a/natten/torch-ext/natten/backends/blackwell_fna.py b/natten/torch-ext/natten/backends/blackwell_fna.py index 184faa0b..5bd6f0ed 100644 --- a/natten/torch-ext/natten/backends/blackwell_fna.py +++ b/natten/torch-ext/natten/backends/blackwell_fna.py @@ -45,7 +45,7 @@ check_cutlass_blackwell_fna_forward_config, ) from ..token_permute import token_permute_operation, token_unpermute_operation -from ..types import ( +from .._types import ( CausalArg1DTypeOrDed, CausalArg2DTypeOrDed, CausalArg3DTypeOrDed, diff --git a/natten/torch-ext/natten/backends/configs/__init__.py b/natten/torch-ext/natten/backends/configs/__init__.py index 732097f2..9c2b75f7 100644 --- a/natten/torch-ext/natten/backends/configs/__init__.py +++ b/natten/torch-ext/natten/backends/configs/__init__.py @@ -62,7 +62,7 @@ get_all_tile_shapes_forward as get_all_flex_fna_forward_configs, get_all_tile_sizes_forward as get_all_flex_fmha_forward_configs, ) -from ...types import ( +from ..._types import ( CutlassBlackwellFmhaBackwardConfigType, CutlassBlackwellFmhaForwardConfigType, CutlassBlackwellFnaBackwardConfigType, diff --git a/natten/torch-ext/natten/backends/configs/cutlass/__init__.py b/natten/torch-ext/natten/backends/configs/cutlass/__init__.py index 7ac0bb72..18760b7b 100644 --- a/natten/torch-ext/natten/backends/configs/cutlass/__init__.py +++ b/natten/torch-ext/natten/backends/configs/cutlass/__init__.py @@ -54,7 +54,7 @@ from ....backends.configs.cutlass.fna_forward_64x64 import ( _FNA_FORWARD_64x64_TILE_SIZES, ) -from ....types import ( +from ...._types import ( CutlassFmhaBackwardConfigType, CutlassFmhaForwardConfigType, CutlassFnaBackwardConfigType, diff --git a/natten/torch-ext/natten/backends/configs/cutlass/backward_knobs.py b/natten/torch-ext/natten/backends/configs/cutlass/backward_knobs.py index da672aaf..dfb1c10a 100644 --- a/natten/torch-ext/natten/backends/configs/cutlass/backward_knobs.py +++ b/natten/torch-ext/natten/backends/configs/cutlass/backward_knobs.py @@ -39,7 +39,7 @@ is_memory_usage_strict, is_memory_usage_unrestricted, ) -from ....types import DimensionType +from ...._types import DimensionType from ....utils.checks import check_dilation_arg, check_input_size_arg from ....utils.tuples import ceil_div_int, ceil_div_tuple diff --git a/natten/torch-ext/natten/backends/configs/cutlass/fna_backward_128x128.py b/natten/torch-ext/natten/backends/configs/cutlass/fna_backward_128x128.py index a7175e21..af6047cc 100644 --- a/natten/torch-ext/natten/backends/configs/cutlass/fna_backward_128x128.py +++ b/natten/torch-ext/natten/backends/configs/cutlass/fna_backward_128x128.py @@ -24,7 +24,7 @@ from typing import Dict, List -from ....types import QKTileShapeType +from ...._types import QKTileShapeType # NOTE: we're excluding tile shapes that include 1 just to # reduce the giant number of configs down to a reasonable diff --git a/natten/torch-ext/natten/backends/configs/cutlass/fna_backward_128x64.py b/natten/torch-ext/natten/backends/configs/cutlass/fna_backward_128x64.py index bfb1e957..6a915e96 100644 --- a/natten/torch-ext/natten/backends/configs/cutlass/fna_backward_128x64.py +++ b/natten/torch-ext/natten/backends/configs/cutlass/fna_backward_128x64.py @@ -24,7 +24,7 @@ from typing import Dict, List -from ....types import CutlassFnaForwardConfigType +from ...._types import CutlassFnaForwardConfigType # NOTE: we're excluding tile shapes that include 1 just to # reduce the giant number of configs down to a reasonable diff --git a/natten/torch-ext/natten/backends/configs/cutlass/fna_backward_64x64.py b/natten/torch-ext/natten/backends/configs/cutlass/fna_backward_64x64.py index 28e9090a..8511bd84 100644 --- a/natten/torch-ext/natten/backends/configs/cutlass/fna_backward_64x64.py +++ b/natten/torch-ext/natten/backends/configs/cutlass/fna_backward_64x64.py @@ -24,7 +24,7 @@ from typing import Dict, List -from ....types import QKTileShapeType +from ...._types import QKTileShapeType # NOTE: we're excluding tile shapes that include 1 just to # reduce the giant number of configs down to a reasonable diff --git a/natten/torch-ext/natten/backends/configs/cutlass/fna_forward_32x128.py b/natten/torch-ext/natten/backends/configs/cutlass/fna_forward_32x128.py index d8e0fd2e..a930aba6 100644 --- a/natten/torch-ext/natten/backends/configs/cutlass/fna_forward_32x128.py +++ b/natten/torch-ext/natten/backends/configs/cutlass/fna_forward_32x128.py @@ -24,7 +24,7 @@ from typing import Dict, List -from ....types import QKTileShapeType +from ...._types import QKTileShapeType # TODO: More combinations are possible for # 2D and 3D (query tile does not have to be smaller diff --git a/natten/torch-ext/natten/backends/configs/cutlass/fna_forward_64x128.py b/natten/torch-ext/natten/backends/configs/cutlass/fna_forward_64x128.py index e6309db4..1bc1fc78 100644 --- a/natten/torch-ext/natten/backends/configs/cutlass/fna_forward_64x128.py +++ b/natten/torch-ext/natten/backends/configs/cutlass/fna_forward_64x128.py @@ -24,7 +24,7 @@ from typing import Dict, List -from ....types import QKTileShapeType +from ...._types import QKTileShapeType # TODO: More combinations are possible for # 2D and 3D (query tile does not have to be smaller diff --git a/natten/torch-ext/natten/backends/configs/cutlass/fna_forward_64x64.py b/natten/torch-ext/natten/backends/configs/cutlass/fna_forward_64x64.py index e70dc716..c321c5b4 100644 --- a/natten/torch-ext/natten/backends/configs/cutlass/fna_forward_64x64.py +++ b/natten/torch-ext/natten/backends/configs/cutlass/fna_forward_64x64.py @@ -24,7 +24,7 @@ from typing import Dict, List -from ....types import QKTileShapeType +from ...._types import QKTileShapeType # TODO: More combinations are possible for # 2D and 3D (query tile does not have to be smaller diff --git a/natten/torch-ext/natten/backends/configs/cutlass_blackwell/__init__.py b/natten/torch-ext/natten/backends/configs/cutlass_blackwell/__init__.py index 165fb92f..4a8c87c6 100644 --- a/natten/torch-ext/natten/backends/configs/cutlass_blackwell/__init__.py +++ b/natten/torch-ext/natten/backends/configs/cutlass_blackwell/__init__.py @@ -26,7 +26,7 @@ import torch # noqa: F401 from torch import Tensor -from ....types import ( +from ...._types import ( CutlassBlackwellFmhaBackwardConfigType, CutlassBlackwellFmhaForwardConfigType, CutlassBlackwellFnaBackwardConfigType, diff --git a/natten/torch-ext/natten/backends/configs/cutlass_hopper/__init__.py b/natten/torch-ext/natten/backends/configs/cutlass_hopper/__init__.py index a22162df..70b2e310 100644 --- a/natten/torch-ext/natten/backends/configs/cutlass_hopper/__init__.py +++ b/natten/torch-ext/natten/backends/configs/cutlass_hopper/__init__.py @@ -26,7 +26,7 @@ import torch # noqa: F401 from torch import Tensor -from ....types import ( +from ...._types import ( CutlassHopperFmhaBackwardConfigType, CutlassHopperFmhaForwardConfigType, CutlassHopperFnaBackwardConfigType, diff --git a/natten/torch-ext/natten/backends/configs/flex/__init__.py b/natten/torch-ext/natten/backends/configs/flex/__init__.py index 0c0e24ee..c5027701 100644 --- a/natten/torch-ext/natten/backends/configs/flex/__init__.py +++ b/natten/torch-ext/natten/backends/configs/flex/__init__.py @@ -26,7 +26,7 @@ import torch # noqa: F401 from torch import Tensor -from ....types import ( +from ...._types import ( DimensionType, FlexFmhaForwardConfigType, FlexFnaForwardConfigType, diff --git a/natten/torch-ext/natten/backends/flex.py b/natten/torch-ext/natten/backends/flex.py index 5ee3b071..0555a7eb 100644 --- a/natten/torch-ext/natten/backends/flex.py +++ b/natten/torch-ext/natten/backends/flex.py @@ -44,7 +44,7 @@ check_flex_fna_forward_config, ) from ..token_permute import token_permute_operation, token_unpermute_operation -from ..types import ( +from .._types import ( CausalArg1DTypeOrDed, CausalArg2DTypeOrDed, CausalArg3DTypeOrDed, diff --git a/natten/torch-ext/natten/backends/fmha.py b/natten/torch-ext/natten/backends/fmha.py index 139f0985..0c6851f8 100644 --- a/natten/torch-ext/natten/backends/fmha.py +++ b/natten/torch-ext/natten/backends/fmha.py @@ -37,7 +37,7 @@ check_cutlass_fmha_backward_config, check_cutlass_fmha_forward_config, ) -from ..types import ( +from .._types import ( CutlassFmhaBackwardConfigType, CutlassFmhaForwardConfigType, NoneType, diff --git a/natten/torch-ext/natten/backends/fna.py b/natten/torch-ext/natten/backends/fna.py index 15be3930..747ecb1f 100644 --- a/natten/torch-ext/natten/backends/fna.py +++ b/natten/torch-ext/natten/backends/fna.py @@ -44,7 +44,7 @@ check_cutlass_fna_backward_config, check_cutlass_fna_forward_config, ) -from ..types import ( +from .._types import ( CausalArg1DTypeOrDed, CausalArg2DTypeOrDed, CausalArg3DTypeOrDed, diff --git a/natten/torch-ext/natten/backends/hopper_fmha.py b/natten/torch-ext/natten/backends/hopper_fmha.py index 4f98a51e..251b50dc 100644 --- a/natten/torch-ext/natten/backends/hopper_fmha.py +++ b/natten/torch-ext/natten/backends/hopper_fmha.py @@ -37,7 +37,7 @@ check_cutlass_hopper_fmha_backward_config, check_cutlass_hopper_fmha_forward_config, ) -from ..types import ( +from .._types import ( CutlassHopperFmhaBackwardConfigType, CutlassHopperFmhaForwardConfigType, KernelSchedule, diff --git a/natten/torch-ext/natten/backends/hopper_fna.py b/natten/torch-ext/natten/backends/hopper_fna.py index dff38175..17bb6d59 100644 --- a/natten/torch-ext/natten/backends/hopper_fna.py +++ b/natten/torch-ext/natten/backends/hopper_fna.py @@ -45,7 +45,7 @@ check_cutlass_hopper_fna_forward_config, ) from ..token_permute import token_permute_operation, token_unpermute_operation -from ..types import ( +from .._types import ( CausalArg1DTypeOrDed, CausalArg2DTypeOrDed, CausalArg3DTypeOrDed, diff --git a/natten/torch-ext/natten/backends/reference.py b/natten/torch-ext/natten/backends/reference.py index c6f72d74..d0f34b41 100644 --- a/natten/torch-ext/natten/backends/reference.py +++ b/natten/torch-ext/natten/backends/reference.py @@ -39,7 +39,7 @@ reference_na3d_backward, reference_na3d_forward, ) -from ..types import ( +from .._types import ( CausalArg1DTypeOrDed, CausalArg2DTypeOrDed, CausalArg3DTypeOrDed, diff --git a/natten/torch-ext/natten/functional.py b/natten/torch-ext/natten/functional.py index d5cd1e42..944a0c2a 100644 --- a/natten/torch-ext/natten/functional.py +++ b/natten/torch-ext/natten/functional.py @@ -38,7 +38,7 @@ flex_fmha, flex_fna_generic, ) -from .types import ( +from ._types import ( CausalArg1DTypeOrDed, CausalArg2DTypeOrDed, CausalArg3DTypeOrDed, diff --git a/natten/torch-ext/natten/modules.py b/natten/torch-ext/natten/modules.py index c3f6a59e..222e384f 100644 --- a/natten/torch-ext/natten/modules.py +++ b/natten/torch-ext/natten/modules.py @@ -26,7 +26,7 @@ from torch import nn, Tensor from .functional import neighborhood_attention_generic -from .types import ( +from ._types import ( CausalArg1DTypeOrDed, CausalArg2DTypeOrDed, CausalArg3DTypeOrDed, diff --git a/natten/torch-ext/natten/token_permute/cutlass_impl.py b/natten/torch-ext/natten/token_permute/cutlass_impl.py index 21694256..49ce678f 100644 --- a/natten/torch-ext/natten/token_permute/cutlass_impl.py +++ b/natten/torch-ext/natten/token_permute/cutlass_impl.py @@ -41,7 +41,7 @@ token_unpermute_2d, token_unpermute_3d, ) -from ..types import DimensionType, NoneType +from .._types import DimensionType, NoneType from ..utils import log from ..utils.device import get_device_cc, is_cuda diff --git a/natten/torch-ext/natten/token_permute/frontend.py b/natten/torch-ext/natten/token_permute/frontend.py index 52d3183f..720aded0 100644 --- a/natten/torch-ext/natten/token_permute/frontend.py +++ b/natten/torch-ext/natten/token_permute/frontend.py @@ -32,7 +32,7 @@ token_unpermute_cutlass, ) from ..token_permute.torch_impl import token_permute_torch, token_unpermute_torch -from ..types import DimensionType +from .._types import DimensionType from ..utils import log from ..utils.tuples import ceil_div_tuple, mul_tuple diff --git a/natten/torch-ext/natten/token_permute/torch_impl.py b/natten/torch-ext/natten/token_permute/torch_impl.py index b0ba3dce..7f06204e 100644 --- a/natten/torch-ext/natten/token_permute/torch_impl.py +++ b/natten/torch-ext/natten/token_permute/torch_impl.py @@ -25,7 +25,7 @@ import torch from torch import Tensor -from ..types import DimensionType +from .._types import DimensionType from ..utils import log from ..utils.environment import is_torch_compiling from ..utils.tuples import ceil_div_tuple, mul_tuple, sub_tuple diff --git a/natten/torch-ext/natten/utils/checks.py b/natten/torch-ext/natten/utils/checks.py index 4c1e339d..f58db50a 100644 --- a/natten/torch-ext/natten/utils/checks.py +++ b/natten/torch-ext/natten/utils/checks.py @@ -28,7 +28,7 @@ import torch # noqa: F401 from torch import Tensor -from ..types import CausalArgType, DimensionType, KernelSchedule, NoneType +from .._types import CausalArgType, DimensionType, KernelSchedule, NoneType from ..utils import log from ..utils.environment import is_torch_compiling from ..utils.tuples import create_causal_arg_from_bool, create_dim_from_int diff --git a/natten/torch-ext/natten/utils/tuples.py b/natten/torch-ext/natten/utils/tuples.py index d1f16de7..abf8a862 100644 --- a/natten/torch-ext/natten/utils/tuples.py +++ b/natten/torch-ext/natten/utils/tuples.py @@ -21,7 +21,7 @@ # ################################################################################################# -from ..types import CausalArgType, DimensionType +from .._types import CausalArgType, DimensionType def ceil_div_int(x: int, y: int) -> int: diff --git a/natten/torch-ext/natten/utils/varlen.py b/natten/torch-ext/natten/utils/varlen.py index 9d4e35d9..8bc541db 100644 --- a/natten/torch-ext/natten/utils/varlen.py +++ b/natten/torch-ext/natten/utils/varlen.py @@ -26,7 +26,7 @@ import torch # noqa: F401 from torch import Tensor -from ..types import NoneType +from .._types import NoneType from ..utils.environment import is_torch_compiling diff --git a/natten/torch-ext/registration.h b/natten/torch-ext/registration.h new file mode 100644 index 00000000..096da4be --- /dev/null +++ b/natten/torch-ext/registration.h @@ -0,0 +1,39 @@ +// Registration macros from vLLM: +// https://github.com/vllm-project/vllm/blob/main/csrc/core/registration.h + +#pragma once + +#include + +#define _CONCAT(A, B) A##B +#define CONCAT(A, B) _CONCAT(A, B) + +#define _STRINGIFY(A) #A +#define STRINGIFY(A) _STRINGIFY(A) + +// A version of the TORCH_LIBRARY macro that expands the NAME, i.e. so NAME +// could be a macro instead of a literal token. +#define TORCH_LIBRARY_EXPAND(NAME, MODULE) TORCH_LIBRARY(NAME, MODULE) + +// A version of the STABLE_TORCH_LIBRARY macro that expands the NAME, i.e. so NAME +// could be a macro instead of a literal token. +#define STABLE_TORCH_LIBRARY_EXPAND(NAME, MODULE) STABLE_TORCH_LIBRARY(NAME, MODULE) + +// A version of the TORCH_LIBRARY_IMPL macro that expands the NAME, i.e. so NAME +// could be a macro instead of a literal token. +#define TORCH_LIBRARY_IMPL_EXPAND(NAME, DEVICE, MODULE) \ + TORCH_LIBRARY_IMPL(NAME, DEVICE, MODULE) + +// A version of the STABLE_TORCH_LIBRARY_IMPL macro that expands the NAME, i.e. so NAME +// could be a macro instead of a literal token. +#define STABLE_TORCH_LIBRARY_IMPL_EXPAND(NAME, DEVICE, MODULE) \ + STABLE_TORCH_LIBRARY_IMPL(NAME, DEVICE, MODULE) + +// REGISTER_EXTENSION allows the shared library to be loaded and initialized +// via python's import statement. +#define REGISTER_EXTENSION(NAME) \ + PyMODINIT_FUNC CONCAT(PyInit_, NAME)() { \ + static struct PyModuleDef module = {PyModuleDef_HEAD_INIT, \ + STRINGIFY(NAME), nullptr, 0, nullptr}; \ + return PyModule_Create(&module); \ + } From 38c1ffee6a63595a4c680ba7b6315cfaa33de1c0 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Fri, 17 Jul 2026 08:09:20 +0530 Subject: [PATCH 6/8] natten: remove accidentally committed create-pyproject artifacts These files are generated per-build by kernel-builder create-pyproject and were committed by mistake in the previous commit. Co-Authored-By: Claude Fable 5 --- natten/CMakeLists.txt | 662 ----------------------------- natten/cmake/_ops.py.in | 9 - natten/cmake/add_build_metadata.py | 56 --- natten/cmake/build-variants.cmake | 328 -------------- natten/cmake/compile-metal.cmake | 104 ----- natten/cmake/get_gpu_lang.cmake | 9 - natten/cmake/get_gpu_lang.py | 20 - natten/cmake/hipify.py | 76 ---- natten/cmake/kernel.cmake | 317 -------------- natten/cmake/metallib_to_header.py | 73 ---- natten/cmake/utils.cmake | 611 -------------------------- natten/compat.py | 26 -- natten/metadata-cann.json | 22 - natten/metadata-cpu.json | 22 - natten/metadata-cuda.json | 22 - natten/metadata-metal.json | 22 - natten/metadata-neuron.json | 22 - natten/metadata-rocm.json | 22 - natten/metadata-xpu.json | 22 - natten/pyproject.toml | 23 - natten/setup.py | 200 --------- natten/torch-ext/registration.h | 39 -- 22 files changed, 2707 deletions(-) delete mode 100644 natten/CMakeLists.txt delete mode 100644 natten/cmake/_ops.py.in delete mode 100644 natten/cmake/add_build_metadata.py delete mode 100644 natten/cmake/build-variants.cmake delete mode 100644 natten/cmake/compile-metal.cmake delete mode 100644 natten/cmake/get_gpu_lang.cmake delete mode 100644 natten/cmake/get_gpu_lang.py delete mode 100644 natten/cmake/hipify.py delete mode 100644 natten/cmake/kernel.cmake delete mode 100644 natten/cmake/metallib_to_header.py delete mode 100644 natten/cmake/utils.cmake delete mode 100644 natten/compat.py delete mode 100644 natten/metadata-cann.json delete mode 100644 natten/metadata-cpu.json delete mode 100644 natten/metadata-cuda.json delete mode 100644 natten/metadata-metal.json delete mode 100644 natten/metadata-neuron.json delete mode 100644 natten/metadata-rocm.json delete mode 100644 natten/metadata-xpu.json delete mode 100644 natten/pyproject.toml delete mode 100644 natten/setup.py delete mode 100644 natten/torch-ext/registration.h diff --git a/natten/CMakeLists.txt b/natten/CMakeLists.txt deleted file mode 100644 index cbf24aec..00000000 --- a/natten/CMakeLists.txt +++ /dev/null @@ -1,662 +0,0 @@ -cmake_minimum_required(VERSION 3.26) - -# Set Intel SYCL compiler before project() call -find_program(ICX_COMPILER icx) -find_program(ICPX_COMPILER icpx) - -if(ICX_COMPILER OR ICPX_COMPILER) - set(CMAKE_C_COMPILER ${ICX_COMPILER}) - - if(WIN32) - set(CMAKE_CXX_COMPILER ${ICX_COMPILER}) - else() - set(CMAKE_CXX_COMPILER ${ICPX_COMPILER}) - endif() -endif() - -project( LANGUAGES CXX) - -install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS) - -include(FetchContent) -file(MAKE_DIRECTORY ${FETCHCONTENT_BASE_DIR}) # Ensure the directory exists -message(STATUS "FetchContent base directory: ${FETCHCONTENT_BASE_DIR}") - -set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1200;gfx1201") - -# Make Torch CMake machinery happy. Whatever we use does not matter, since -# we set the arches per-file later anyway. -set(ENV{PYTORCH_ROCM_ARCH} "${HIP_SUPPORTED_ARCHS}") - -include(${CMAKE_CURRENT_LIST_DIR}/cmake/utils.cmake) -include(${CMAKE_CURRENT_LIST_DIR}/cmake/kernel.cmake) -include(${CMAKE_CURRENT_LIST_DIR}/cmake/get_gpu_lang.cmake) - -if(DEFINED Python3_EXECUTABLE) - # Allow passing through the interpreter (e.g. from setup.py). - find_package(Python3 COMPONENTS Development Development.SABIModule Interpreter) - if (NOT Python3_FOUND) - message(FATAL_ERROR "Unable to find python matching: ${EXECUTABLE}.") - endif() -else() - find_package(Python3 REQUIRED COMPONENTS Development Development.SABIModule Interpreter) -endif() - -get_gpu_lang(DETECTED_GPU_LANG) -set(GPU_LANG "${DETECTED_GPU_LANG}" CACHE STRING "GPU language") -gpu_lang_to_backend(BACKEND "${GPU_LANG}") -message(STATUS "Using backend: ${BACKEND}, GPU language: ${GPU_LANG}") - -set(OPS_NAME "_natten_${BACKEND}_28fa1dd") - -append_cmake_prefix_path("torch" "torch.utils.cmake_prefix_path") - -find_package(Torch REQUIRED) - -run_python(TORCH_VERSION "import torch; print(torch.__version__.split('+')[0])" "Failed to get Torch version") - - - - -option(BUILD_ALL_SUPPORTED_ARCHS "Build all supported architectures" off) - -if(DEFINED CMAKE_CUDA_COMPILER_VERSION AND - CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0) - set(CUDA_DEFAULT_KERNEL_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;11.0;12.0;12.1+PTX") -elseif(DEFINED CMAKE_CUDA_COMPILER_VERSION AND - CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.8) - set(CUDA_DEFAULT_KERNEL_ARCHS "7.0;7.2;7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.1;12.0+PTX") -else() - set(CUDA_DEFAULT_KERNEL_ARCHS "7.0;7.2;7.5;8.0;8.6;8.7;8.9;9.0+PTX") -endif() - -# Basic checks for each GPU language. -if(GPU_LANG STREQUAL "CUDA") - if(NOT CUDA_FOUND) - message(FATAL_ERROR "GPU language is set to CUDA, but cannot find CUDA toolkit") - endif() - - - - # This clears out -gencode arguments from `CMAKE_CUDA_FLAGS`, which we need - # to set our own set of capabilities. - clear_gencode_flags() - - # Get the capabilities without +PTX suffixes, so that we can use them as - # the target archs in the loose intersection with a kernel's capabilities. - cuda_remove_ptx_suffixes(CUDA_ARCHS "${CUDA_DEFAULT_KERNEL_ARCHS}") - message(STATUS "CUDA base archs used for intersection with kernel archs: ${CUDA_ARCHS}") - - if(BUILD_ALL_SUPPORTED_ARCHS) - set(CUDA_KERNEL_ARCHS "${CUDA_DEFAULT_KERNEL_ARCHS}") - else() - try_run_python(CUDA_KERNEL_ARCHS SUCCESS "import torch; cc=torch.cuda.get_device_capability(); print(f\"{cc[0]}.{cc[1]}\")" "Failed to get CUDA capability") - if(NOT SUCCESS) - message(WARNING "Failed to detect CUDA capability, using default capabilities.") - set(CUDA_KERNEL_ARCHS "${CUDA_DEFAULT_KERNEL_ARCHS}") - endif() - endif() - - message(STATUS "Default CUDA kernel architectures: ${CUDA_KERNEL_ARCHS}") - - if(NVCC_THREADS AND GPU_LANG STREQUAL "CUDA") - message(STATUS "Using nvcc with: -threads=${NVCC_THREADS}") - list(APPEND GPU_FLAGS "--threads=${NVCC_THREADS}") - endif() - - # TODO: deprecate one of these settings. - add_compile_definitions(USE_CUDA=1) - add_compile_definitions(CUDA_KERNEL) -elseif(GPU_LANG STREQUAL "HIP") - if(NOT HIP_FOUND AND NOT PYTORCH_FOUND_HIP) - message(FATAL_ERROR "GPU language is set to HIP, but cannot find ROCm toolkit") - endif() - - # Importing torch recognizes and sets up some HIP/ROCm configuration but does - # not let cmake recognize .hip files. In order to get cmake to understand the - # .hip extension automatically, HIP must be enabled explicitly. - enable_language(HIP) - - set(ROCM_ARCHS ${HIP_SUPPORTED_ARCHS}) - message(STATUS "ROCM supported target architectures: ${ROCM_ARCHS}") - - # TODO: deprecate one of these settings. - add_compile_definitions(USE_ROCM=1) - add_compile_definitions(ROCM_KERNEL) -elseif(GPU_LANG STREQUAL "CPU") - add_compile_definitions(CPU_KERNEL) - set(CMAKE_OSX_DEPLOYMENT_TARGET "15.0" CACHE STRING "Minimum macOS deployment version") -elseif(GPU_LANG STREQUAL "METAL") - set(CMAKE_OSX_DEPLOYMENT_TARGET "26.0" CACHE STRING "Minimum macOS deployment version") - enable_language(C OBJC OBJCXX) - - add_compile_definitions(METAL_KERNEL) - - # Initialize lists for Metal shader sources and their include directories - set(ALL_METAL_SOURCES) - set(METAL_INCLUDE_DIRS) -elseif(GPU_LANG STREQUAL "SYCL") - if(NOT ICX_COMPILER AND NOT ICPX_COMPILER) - message(FATAL_ERROR "Intel SYCL C++ compiler (icpx) and/or C compiler (icx) not found. Please install Intel oneAPI toolkit.") - endif() - - execute_process( - COMMAND ${ICPX_COMPILER} --version - OUTPUT_VARIABLE ICPX_VERSION_OUTPUT - OUTPUT_STRIP_TRAILING_WHITESPACE - ) - string(REGEX MATCH "[0-9]+\\.[0-9]+" DPCPP_VERSION "${ICPX_VERSION_OUTPUT}") - set(DPCPP_VERSION "${DPCPP_VERSION}" CACHE STRING "DPCPP major.minor version") - - # On Windows, use icx (MSVC-compatible) for C++ to work with Ninja generator - # On Linux, use icpx (GNU-compatible) for C++ - if(WIN32) - message(STATUS "Using Intel SYCL C++ compiler: ${ICX_COMPILER} and C compiler: ${ICX_COMPILER} Version: ${DPCPP_VERSION} (Windows MSVC-compatible mode)") - else() - message(STATUS "Using Intel SYCL C++ compiler: ${ICPX_COMPILER} and C compiler: ${ICX_COMPILER} Version: ${DPCPP_VERSION}") - endif() - - - set(sycl_link_flags "-Wl,-z,noexecstack;-fsycl;--offload-compress;-fsycl-targets=spir64_gen,spir64;-Xs;-device pvc,xe-lpg,ats-m150 -options ' -cl-intel-enable-auto-large-GRF-mode -cl-poison-unsupported-fp64-kernels -cl-intel-greater-than-4GB-buffer-required';") - set(sycl_flags "-fPIC;-fsycl;-fhonor-nans;-fhonor-infinities;-fno-associative-math;-fno-approx-func;-fno-sycl-instrument-device-code;--offload-compress;-fsycl-targets=spir64_gen,spir64;") - set(GPU_FLAGS "${sycl_flags}") - - - add_compile_definitions(XPU_KERNEL) - add_compile_definitions(USE_XPU) -else() - message(FATAL_ERROR "Unsupported GPU language: ${GPU_LANG}") -endif() - -# Initialize SRC list for kernel and binding sources -set(SRC "") - -include(${CMAKE_CURRENT_LIST_DIR}/cmake/build-variants.cmake) - -# Generate build variant name. -if(GPU_LANG STREQUAL "CUDA") - generate_build_name(BUILD_VARIANT_NAME "${TORCH_VERSION}" "cuda" "${CUDA_VERSION}" ${_STABLE_ABI_ARG}) -elseif(GPU_LANG STREQUAL "HIP") - run_python(ROCM_VERSION "import torch.version; print(torch.version.hip.split('.')[0] + '.' + torch.version.hip.split('.')[1])" "Failed to get ROCm version") - generate_build_name(BUILD_VARIANT_NAME "${TORCH_VERSION}" "rocm" "${ROCM_VERSION}" ${_STABLE_ABI_ARG}) -elseif(GPU_LANG STREQUAL "SYCL") - generate_build_name(BUILD_VARIANT_NAME "${TORCH_VERSION}" "xpu" "${DPCPP_VERSION}" ${_STABLE_ABI_ARG}) -elseif(GPU_LANG STREQUAL "METAL") - generate_build_name(BUILD_VARIANT_NAME "${TORCH_VERSION}" "metal" "" ${_STABLE_ABI_ARG}) -elseif(GPU_LANG STREQUAL "CPU") - generate_build_name(BUILD_VARIANT_NAME "${TORCH_VERSION}" "cpu" "" ${_STABLE_ABI_ARG}) -else() - message(FATAL_ERROR "Cannot generate build name for unknown GPU_LANG: ${GPU_LANG}") -endif() - -configure_file( - ${CMAKE_CURRENT_LIST_DIR}/cmake/_ops.py.in - ${CMAKE_CURRENT_SOURCE_DIR}/torch-ext/natten/_ops.py - @ONLY -) - -if(GPU_LANG STREQUAL "CUDA") - message(STATUS "Including CUTLASS dependency") - -find_package(NvidiaCutlass) - -if (NOT NvidiaCutlass_FOUND) - set(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library") - -# Set CUTLASS_REVISION manually -- its revision detection doesn't work in this case. - set(CUTLASS_REVISION "v4.5.2" CACHE STRING "CUTLASS revision to use") - - -# Use the specified CUTLASS source directory for compilation if CUTLASS_SRC_DIR is provided - if (DEFINED ENV{CUTLASS_SRC_DIR}) - set(CUTLASS_SRC_DIR $ENV{CUTLASS_SRC_DIR}) - endif() - - if(CUTLASS_SRC_DIR) - if(NOT IS_ABSOLUTE CUTLASS_SRC_DIR) - get_filename_component(CUTLASS_SRC_DIR "${CUTLASS_SRC_DIR}" ABSOLUTE) - endif() - message(STATUS "The CUTLASS_SRC_DIR is set, using ${CUTLASS_SRC_DIR} for compilation") - FetchContent_Declare(cutlass SOURCE_DIR ${CUTLASS_SRC_DIR}) - else() - FetchContent_Declare( - cutlass - GIT_REPOSITORY https://github.com/nvidia/cutlass.git - GIT_TAG ${CUTLASS_REVISION} - GIT_PROGRESS TRUE - - # Speed up CUTLASS download by retrieving only the specified GIT_TAG instead of the history. - # Important: If GIT_SHALLOW is enabled then GIT_TAG works only with branch names and tags. - # So if the GIT_TAG above is updated to a commit hash, GIT_SHALLOW must be set to FALSE - GIT_SHALLOW TRUE - ) - endif() - FetchContent_MakeAvailable(cutlass) - - include_directories(${CUTLASS_INCLUDE_DIR}) -else() - message(STATUS "Using system cutlass with version: ${NvidiaCutlass_VERSION}") -endif(NOT NvidiaCutlass_FOUND) - -endif(GPU_LANG STREQUAL "CUDA") -if(GPU_LANG STREQUAL "CUDA") - get_torch_gpu_compiler_flags(TORCH_GPU_FLAGS ${GPU_LANG}) - list(APPEND GPU_FLAGS ${TORCH_GPU_FLAGS}) -endif() - -set(TORCH_natten_SRC - torch-ext/torch_binding.cpp -) - -# TODO: check if CLion support this: -# https://youtrack.jetbrains.com/issue/CPP-16510/CLion-does-not-handle-per-file-include-directories -set_source_files_properties( - ${TORCH_natten_SRC} - PROPERTIES INCLUDE_DIRECTORIES "${CMAKE_SOURCE_DIR}/csrc/include") - - -list(APPEND SRC "${TORCH_natten_SRC}") -if(GPU_LANG STREQUAL "CUDA") -cuda_kernel_component(SRC - NAME natten_blackwell - SOURCES "csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_cm.h" -"csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_dtype.h" -"csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_head_dim.h" -"csrc/autogen/include/natten_autogen/cuda/blackwell_fna/dispatch_tile_shape.h" -"csrc/autogen/include/natten_autogen/cuda/blackwell_fna/interface.h" -"csrc/autogen/include/natten_autogen/cuda/blackwell_fna/kernels.h" -"csrc/autogen/src/cuda/blackwell_fna/source_0.cu" -"csrc/autogen/src/cuda/blackwell_fna/source_1.cu" -"csrc/autogen/src/cuda/blackwell_fna/source_10.cu" -"csrc/autogen/src/cuda/blackwell_fna/source_11.cu" -"csrc/autogen/src/cuda/blackwell_fna/source_12.cu" -"csrc/autogen/src/cuda/blackwell_fna/source_13.cu" -"csrc/autogen/src/cuda/blackwell_fna/source_14.cu" -"csrc/autogen/src/cuda/blackwell_fna/source_15.cu" -"csrc/autogen/src/cuda/blackwell_fna/source_16.cu" -"csrc/autogen/src/cuda/blackwell_fna/source_17.cu" -"csrc/autogen/src/cuda/blackwell_fna/source_18.cu" -"csrc/autogen/src/cuda/blackwell_fna/source_19.cu" -"csrc/autogen/src/cuda/blackwell_fna/source_2.cu" -"csrc/autogen/src/cuda/blackwell_fna/source_20.cu" -"csrc/autogen/src/cuda/blackwell_fna/source_21.cu" -"csrc/autogen/src/cuda/blackwell_fna/source_22.cu" -"csrc/autogen/src/cuda/blackwell_fna/source_23.cu" -"csrc/autogen/src/cuda/blackwell_fna/source_24.cu" -"csrc/autogen/src/cuda/blackwell_fna/source_25.cu" -"csrc/autogen/src/cuda/blackwell_fna/source_26.cu" -"csrc/autogen/src/cuda/blackwell_fna/source_27.cu" -"csrc/autogen/src/cuda/blackwell_fna/source_3.cu" -"csrc/autogen/src/cuda/blackwell_fna/source_4.cu" -"csrc/autogen/src/cuda/blackwell_fna/source_5.cu" -"csrc/autogen/src/cuda/blackwell_fna/source_6.cu" -"csrc/autogen/src/cuda/blackwell_fna/source_7.cu" -"csrc/autogen/src/cuda/blackwell_fna/source_8.cu" -"csrc/autogen/src/cuda/blackwell_fna/source_9.cu" -"csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_cm.h" -"csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_dtype.h" -"csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_head_dim.h" -"csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/dispatch_tile_shape.h" -"csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/interface.h" -"csrc/autogen/include/natten_autogen/cuda/blackwell_fna_bwd/kernels.h" -"csrc/autogen/src/cuda/blackwell_fna_bwd/source_0.cu" -"csrc/autogen/src/cuda/blackwell_fna_bwd/source_1.cu" -"csrc/autogen/src/cuda/blackwell_fna_bwd/source_10.cu" -"csrc/autogen/src/cuda/blackwell_fna_bwd/source_11.cu" -"csrc/autogen/src/cuda/blackwell_fna_bwd/source_12.cu" -"csrc/autogen/src/cuda/blackwell_fna_bwd/source_13.cu" -"csrc/autogen/src/cuda/blackwell_fna_bwd/source_2.cu" -"csrc/autogen/src/cuda/blackwell_fna_bwd/source_3.cu" -"csrc/autogen/src/cuda/blackwell_fna_bwd/source_4.cu" -"csrc/autogen/src/cuda/blackwell_fna_bwd/source_5.cu" -"csrc/autogen/src/cuda/blackwell_fna_bwd/source_6.cu" -"csrc/autogen/src/cuda/blackwell_fna_bwd/source_7.cu" -"csrc/autogen/src/cuda/blackwell_fna_bwd/source_8.cu" -"csrc/autogen/src/cuda/blackwell_fna_bwd/source_9.cu" -"csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/dispatch_head_dim.h" -"csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/dispatch_tile_size.h" -"csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/interface.h" -"csrc/autogen/include/natten_autogen/cuda/blackwell_fmha/kernels.h" -"csrc/autogen/src/cuda/blackwell_fmha/source_0.cu" -"csrc/autogen/src/cuda/blackwell_fmha/source_1.cu" -"csrc/autogen/src/cuda/blackwell_fmha/source_2.cu" -"csrc/autogen/src/cuda/blackwell_fmha/source_3.cu" -"csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/dispatch_head_dim.h" -"csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/dispatch_tile_size.h" -"csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/interface.h" -"csrc/autogen/include/natten_autogen/cuda/blackwell_fmha_bwd/kernels.h" -"csrc/autogen/src/cuda/blackwell_fmha_bwd/source_0.cu" -"csrc/autogen/src/cuda/blackwell_fmha_bwd/source_1.cu" -"csrc/autogen/src/cuda/blackwell_fmha_bwd/source_2.cu" -"csrc/autogen/src/cuda/blackwell_fmha_bwd/source_3.cu" - CUDA_MINVER 12.8 INCLUDES "${CMAKE_SOURCE_DIR}/csrc/include;${CMAKE_SOURCE_DIR}/csrc/autogen/include" CUDA_CAPABILITIES 10.0a CUDA_FLAGS "-O3;-std=c++17;--expt-relaxed-constexpr;--extended-lambda;--use_fast_math;--ftemplate-backtrace-limit=0;-Xfatbin=-compress-all;-DNDEBUG;-DNATTEN_WITH_CUTLASS;-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1;-DNATTEN_WITH_BLACKWELL_FNA=1" ) -endif() -if(GPU_LANG STREQUAL "CUDA") -cuda_kernel_component(SRC - NAME natten - SOURCES "csrc/include/natten/blackwell_fmha.h" -"csrc/include/natten/blackwell_fna.h" -"csrc/include/natten/compute_delta.h" -"csrc/include/natten/cuda/fmha/epilogue/epilogue_pipelined.h" -"csrc/include/natten/cuda/fmha/epilogue/epilogue_rescale_output.h" -"csrc/include/natten/cuda/fmha/epilogue/epilogue_thread_apply_logsumexp.h" -"csrc/include/natten/cuda/fmha/fmha_backward.cuh" -"csrc/include/natten/cuda/fmha/fmha_forward.cuh" -"csrc/include/natten/cuda/fmha/gemm/custom_mma.h" -"csrc/include/natten/cuda/fmha/gemm/custom_mma_base.h" -"csrc/include/natten/cuda/fmha/gemm/custom_mma_multistage.h" -"csrc/include/natten/cuda/fmha/gemm/custom_mma_pipelined.h" -"csrc/include/natten/cuda/fmha/gemm/find_default_mma.h" -"csrc/include/natten/cuda/fmha/gemm/mma_accum_lambda_iterator.h" -"csrc/include/natten/cuda/fmha/gemm/mma_from_smem.h" -"csrc/include/natten/cuda/fmha/gemm_kernel_utils.h" -"csrc/include/natten/cuda/fmha/iterators/default_warp_iterator_from_smem.h" -"csrc/include/natten/cuda/fmha/iterators/epilogue_predicated_tile_iterator.h" -"csrc/include/natten/cuda/fmha/iterators/make_residual_last.h" -"csrc/include/natten/cuda/fmha/iterators/predicated_tile_access_iterator_residual_last.h" -"csrc/include/natten/cuda/fmha/iterators/predicated_tile_iterator_residual_last.h" -"csrc/include/natten/cuda/fmha/iterators/transpose_warp_iterator.h" -"csrc/include/natten/cuda/fmha/iterators/warp_iterator_from_smem.h" -"csrc/include/natten/cuda/fmha/kernel_backward.h" -"csrc/include/natten/cuda/fmha/kernel_forward.h" -"csrc/include/natten/cuda/fmha/transform/tile_smem_loader.h" -"csrc/include/natten/cuda/fmha_blackwell/collective/fmha_common.hpp" -"csrc/include/natten/cuda/fmha_blackwell/collective/fmha_fusion.hpp" -"csrc/include/natten/cuda/fmha_blackwell/collective/sm100_fmha_fwd_epilogue_tma_warpspecialized.hpp" -"csrc/include/natten/cuda/fmha_blackwell/collective/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp" -"csrc/include/natten/cuda/fmha_blackwell/collective/sm100_fmha_load_tma_warpspecialized.hpp" -"csrc/include/natten/cuda/fmha_blackwell/common/pow_2.hpp" -"csrc/include/natten/cuda/fmha_blackwell/device/fmha_bwd_sm100.hpp" -"csrc/include/natten/cuda/fmha_blackwell/device/fmha_sm100.hpp" -"csrc/include/natten/cuda/fmha_blackwell/fmha_backward.cuh" -"csrc/include/natten/cuda/fmha_blackwell/fmha_forward.cuh" -"csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_causal_tile_scheduler.hpp" -"csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_kernel_bwd_convert.hpp" -"csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_kernel_bwd_sum_OdO.hpp" -"csrc/include/natten/cuda/fmha_blackwell/kernel/fmha_tile_scheduler.hpp" -"csrc/include/natten/cuda/fmha_blackwell/kernel/sm100_fmha_bwd_kernel_tma_warpspecialized.hpp" -"csrc/include/natten/cuda/fmha_blackwell/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp" -"csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_bwd_tma_warpspecialized.hpp" -"csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_load.hpp" -"csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_softmax.hpp" -"csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_tma.hpp" -"csrc/include/natten/cuda/fmha_hopper/collective/fmha_collective_tma_warpspecialized.hpp" -"csrc/include/natten/cuda/fmha_hopper/collective/fmha_common.hpp" -"csrc/include/natten/cuda/fmha_hopper/collective/fmha_epilogue.hpp" -"csrc/include/natten/cuda/fmha_hopper/collective/fmha_epilogue_bwd.hpp" -"csrc/include/natten/cuda/fmha_hopper/collective/fmha_fusion.hpp" -"csrc/include/natten/cuda/fmha_hopper/collective/fmha_varlen.hpp" -"csrc/include/natten/cuda/fmha_hopper/device/fmha_bwd_sm90.hpp" -"csrc/include/natten/cuda/fmha_hopper/device/fmha_sm90.hpp" -"csrc/include/natten/cuda/fmha_hopper/fmha_backward.cuh" -"csrc/include/natten/cuda/fmha_hopper/fmha_forward.cuh" -"csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_builder.hpp" -"csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_bwd_convert.hpp" -"csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_bwd_sum_OdO.hpp" -"csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_tma.hpp" -"csrc/include/natten/cuda/fmha_hopper/kernel/fmha_kernel_tma_warpspecialized.hpp" -"csrc/include/natten/cuda/fmha_hopper/kernel/fmha_options.hpp" -"csrc/include/natten/cuda/fmha_hopper/kernel/fmha_tile_scheduler.hpp" -"csrc/include/natten/cuda/fna/epilogue/epilogue_pipelined.h" -"csrc/include/natten/cuda/fna/epilogue/epilogue_rescale_output.h" -"csrc/include/natten/cuda/fna/epilogue/epilogue_thread_apply_logsumexp.h" -"csrc/include/natten/cuda/fna/epilogue/predicated_tile_iterator.h" -"csrc/include/natten/cuda/fna/epilogue/predicated_tile_iterator_params.h" -"csrc/include/natten/cuda/fna/fna_backward.cuh" -"csrc/include/natten/cuda/fna/fna_forward.cuh" -"csrc/include/natten/cuda/fna/gemm/custom_mma.h" -"csrc/include/natten/cuda/fna/gemm/custom_mma_base.h" -"csrc/include/natten/cuda/fna/gemm/custom_mma_multistage.h" -"csrc/include/natten/cuda/fna/gemm/custom_mma_pipelined.h" -"csrc/include/natten/cuda/fna/gemm/find_default_mma.h" -"csrc/include/natten/cuda/fna/gemm/mma_accum_lambda_iterator.h" -"csrc/include/natten/cuda/fna/gemm/mma_from_smem.h" -"csrc/include/natten/cuda/fna/gemm/replace_mma_iterators.h" -"csrc/include/natten/cuda/fna/gemm_kernel_utils.h" -"csrc/include/natten/cuda/fna/iterators/default_warp_iterator_from_smem.h" -"csrc/include/natten/cuda/fna/iterators/epilogue_predicated_tile_iterator.h" -"csrc/include/natten/cuda/fna/iterators/make_residual_last.h" -"csrc/include/natten/cuda/fna/iterators/predicated_tile_access_iterator.h" -"csrc/include/natten/cuda/fna/iterators/predicated_tile_access_iterator_residual_last.h" -"csrc/include/natten/cuda/fna/iterators/predicated_tile_iterator.h" -"csrc/include/natten/cuda/fna/iterators/predicated_tile_iterator_residual_last.h" -"csrc/include/natten/cuda/fna/iterators/transpose_warp_iterator.h" -"csrc/include/natten/cuda/fna/iterators/warp_iterator_from_smem.h" -"csrc/include/natten/cuda/fna/kernel_backward.h" -"csrc/include/natten/cuda/fna/kernel_forward.h" -"csrc/include/natten/cuda/fna/na_utils.cuh" -"csrc/include/natten/cuda/fna_blackwell/collective/fna_common.hpp" -"csrc/include/natten/cuda/fna_blackwell/collective/fna_fusion.hpp" -"csrc/include/natten/cuda/fna_blackwell/collective/fna_fusion_bwd.hpp" -"csrc/include/natten/cuda/fna_blackwell/collective/sm100_fna_fwd_epilogue_tma_warpspecialized.hpp" -"csrc/include/natten/cuda/fna_blackwell/collective/sm100_fna_fwd_mainloop_tma_warpspecialized.hpp" -"csrc/include/natten/cuda/fna_blackwell/collective/sm100_fna_load_tma_warpspecialized.hpp" -"csrc/include/natten/cuda/fna_blackwell/common/pow_2.hpp" -"csrc/include/natten/cuda/fna_blackwell/device/fna_bwd_sm100.hpp" -"csrc/include/natten/cuda/fna_blackwell/device/fna_sm100.hpp" -"csrc/include/natten/cuda/fna_blackwell/fna_backward.cuh" -"csrc/include/natten/cuda/fna_blackwell/fna_forward.cuh" -"csrc/include/natten/cuda/fna_blackwell/kernel/sm100_fna_bwd_kernel_tma_warpspecialized.hpp" -"csrc/include/natten/cuda/fna_blackwell/kernel/sm100_fna_fwd_kernel_tma_warpspecialized.hpp" -"csrc/include/natten/cuda/fna_hopper/collective/fna_collective_bwd_tma_warpspecialized.hpp" -"csrc/include/natten/cuda/fna_hopper/collective/fna_collective_load.hpp" -"csrc/include/natten/cuda/fna_hopper/collective/fna_collective_softmax.hpp" -"csrc/include/natten/cuda/fna_hopper/collective/fna_collective_tma.hpp" -"csrc/include/natten/cuda/fna_hopper/collective/fna_collective_tma_warpspecialized.hpp" -"csrc/include/natten/cuda/fna_hopper/collective/fna_common.hpp" -"csrc/include/natten/cuda/fna_hopper/collective/fna_fusion.hpp" -"csrc/include/natten/cuda/fna_hopper/collective/fna_fusion_bwd.hpp" -"csrc/include/natten/cuda/fna_hopper/device/fna_bwd_sm90.hpp" -"csrc/include/natten/cuda/fna_hopper/device/fna_sm90.hpp" -"csrc/include/natten/cuda/fna_hopper/fna_backward.cuh" -"csrc/include/natten/cuda/fna_hopper/fna_forward.cuh" -"csrc/include/natten/cuda/fna_hopper/kernel/fna_kernel_builder.hpp" -"csrc/include/natten/cuda/hopper_fmha_fna.h" -"csrc/include/natten/cuda/reduction/compute_delta.cuh" -"csrc/include/natten/cuda/reduction/fmha_kernel_bwd_sum_OdO.hpp" -"csrc/include/natten/cuda/reference/fna_reference_backward.hpp" -"csrc/include/natten/cuda/reference/fna_reference_forward.hpp" -"csrc/include/natten/cuda/reference/mask.hpp" -"csrc/include/natten/cuda/reference/utils.hpp" -"csrc/include/natten/cuda/tokperm/layouts.hpp" -"csrc/include/natten/cuda/tokperm/token_permute_kernel.cuh" -"csrc/include/natten/cuda/tokperm/tokperm.hpp" -"csrc/include/natten/cuda/tokperm/utils/permute.cuh" -"csrc/include/natten/cuda/tokperm/utils/stride.cuh" -"csrc/include/natten/cuda/tokperm/utils/tuple.cuh" -"csrc/include/natten/cuda/utils/cuda.h" -"csrc/include/natten/cuda/utils/cutlass.cuh" -"csrc/include/natten/cuda/utils/generic_cutlass_device.hpp" -"csrc/include/natten/fmha.h" -"csrc/include/natten/fna.h" -"csrc/include/natten/helpers.h" -"csrc/include/natten/hopper_fmha.h" -"csrc/include/natten/hopper_fna.h" -"csrc/include/natten/natten.h" -"csrc/include/natten/reference.h" -"csrc/include/natten/token_permute.h" -"csrc/src/blackwell_fmha.cu" -"csrc/src/blackwell_fna_backward.cu" -"csrc/src/blackwell_fna_forward.cu" -"csrc/src/compute_delta.cu" -"csrc/src/fmha.cu" -"csrc/src/fna_backward.cu" -"csrc/src/fna_forward.cu" -"csrc/src/hopper_fmha.cu" -"csrc/src/hopper_fna_backward.cu" -"csrc/src/hopper_fna_forward.cu" -"csrc/src/reference_backward.cu" -"csrc/src/reference_forward.cu" -"csrc/src/token_permute.cu" -"csrc/src/token_unpermute.cu" -"csrc/autogen/include/natten_autogen/cuda/reference/dispatch_cm.h" -"csrc/autogen/include/natten_autogen/cuda/reference/dispatch_dtype.h" -"csrc/autogen/include/natten_autogen/cuda/reference/interface.h" -"csrc/autogen/include/natten_autogen/cuda/reference/kernels.h" -"csrc/autogen/src/cuda/reference/source_0.cu" -"csrc/autogen/src/cuda/reference/source_1.cu" -"csrc/autogen/include/natten_autogen/cuda/fna/dispatch_cm.h" -"csrc/autogen/include/natten_autogen/cuda/fna/dispatch_device.h" -"csrc/autogen/include/natten_autogen/cuda/fna/dispatch_dtype.h" -"csrc/autogen/include/natten_autogen/cuda/fna/interface.h" -"csrc/autogen/include/natten_autogen/cuda/fna/kernels.h" -"csrc/autogen/src/cuda/fna/source_0.cu" -"csrc/autogen/src/cuda/fna/source_1.cu" -"csrc/autogen/src/cuda/fna/source_10.cu" -"csrc/autogen/src/cuda/fna/source_11.cu" -"csrc/autogen/src/cuda/fna/source_12.cu" -"csrc/autogen/src/cuda/fna/source_13.cu" -"csrc/autogen/src/cuda/fna/source_14.cu" -"csrc/autogen/src/cuda/fna/source_15.cu" -"csrc/autogen/src/cuda/fna/source_16.cu" -"csrc/autogen/src/cuda/fna/source_17.cu" -"csrc/autogen/src/cuda/fna/source_18.cu" -"csrc/autogen/src/cuda/fna/source_19.cu" -"csrc/autogen/src/cuda/fna/source_2.cu" -"csrc/autogen/src/cuda/fna/source_20.cu" -"csrc/autogen/src/cuda/fna/source_21.cu" -"csrc/autogen/src/cuda/fna/source_22.cu" -"csrc/autogen/src/cuda/fna/source_23.cu" -"csrc/autogen/src/cuda/fna/source_24.cu" -"csrc/autogen/src/cuda/fna/source_25.cu" -"csrc/autogen/src/cuda/fna/source_26.cu" -"csrc/autogen/src/cuda/fna/source_27.cu" -"csrc/autogen/src/cuda/fna/source_28.cu" -"csrc/autogen/src/cuda/fna/source_29.cu" -"csrc/autogen/src/cuda/fna/source_3.cu" -"csrc/autogen/src/cuda/fna/source_30.cu" -"csrc/autogen/src/cuda/fna/source_31.cu" -"csrc/autogen/src/cuda/fna/source_32.cu" -"csrc/autogen/src/cuda/fna/source_33.cu" -"csrc/autogen/src/cuda/fna/source_34.cu" -"csrc/autogen/src/cuda/fna/source_35.cu" -"csrc/autogen/src/cuda/fna/source_36.cu" -"csrc/autogen/src/cuda/fna/source_37.cu" -"csrc/autogen/src/cuda/fna/source_38.cu" -"csrc/autogen/src/cuda/fna/source_39.cu" -"csrc/autogen/src/cuda/fna/source_4.cu" -"csrc/autogen/src/cuda/fna/source_40.cu" -"csrc/autogen/src/cuda/fna/source_41.cu" -"csrc/autogen/src/cuda/fna/source_42.cu" -"csrc/autogen/src/cuda/fna/source_43.cu" -"csrc/autogen/src/cuda/fna/source_44.cu" -"csrc/autogen/src/cuda/fna/source_45.cu" -"csrc/autogen/src/cuda/fna/source_46.cu" -"csrc/autogen/src/cuda/fna/source_47.cu" -"csrc/autogen/src/cuda/fna/source_48.cu" -"csrc/autogen/src/cuda/fna/source_49.cu" -"csrc/autogen/src/cuda/fna/source_5.cu" -"csrc/autogen/src/cuda/fna/source_50.cu" -"csrc/autogen/src/cuda/fna/source_51.cu" -"csrc/autogen/src/cuda/fna/source_52.cu" -"csrc/autogen/src/cuda/fna/source_53.cu" -"csrc/autogen/src/cuda/fna/source_54.cu" -"csrc/autogen/src/cuda/fna/source_55.cu" -"csrc/autogen/src/cuda/fna/source_56.cu" -"csrc/autogen/src/cuda/fna/source_57.cu" -"csrc/autogen/src/cuda/fna/source_58.cu" -"csrc/autogen/src/cuda/fna/source_59.cu" -"csrc/autogen/src/cuda/fna/source_6.cu" -"csrc/autogen/src/cuda/fna/source_60.cu" -"csrc/autogen/src/cuda/fna/source_61.cu" -"csrc/autogen/src/cuda/fna/source_62.cu" -"csrc/autogen/src/cuda/fna/source_63.cu" -"csrc/autogen/src/cuda/fna/source_7.cu" -"csrc/autogen/src/cuda/fna/source_8.cu" -"csrc/autogen/src/cuda/fna/source_9.cu" -"csrc/autogen/include/natten_autogen/cuda/fmha/dispatch_dtype.h" -"csrc/autogen/include/natten_autogen/cuda/fmha/interface.h" -"csrc/autogen/include/natten_autogen/cuda/fmha/kernels.h" -"csrc/autogen/src/cuda/fmha/source_0.cu" -"csrc/autogen/src/cuda/fmha/source_1.cu" -"csrc/autogen/src/cuda/fmha/source_2.cu" -"csrc/autogen/src/cuda/fmha/source_3.cu" -"csrc/autogen/src/cuda/fmha/source_4.cu" -"csrc/autogen/src/cuda/fmha/source_5.cu" - INCLUDES "${CMAKE_SOURCE_DIR}/csrc/include;${CMAKE_SOURCE_DIR}/csrc/autogen/include" CUDA_CAPABILITIES 7.0 7.5 8.0 8.6 8.9 9.0 10.0 12.0 CUDA_FLAGS "-O3;-std=c++17;--expt-relaxed-constexpr;--extended-lambda;--use_fast_math;--ftemplate-backtrace-limit=0;-Xfatbin=-compress-all;-DNDEBUG;-DNATTEN_WITH_CUTLASS;-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1" ) -endif() -if(GPU_LANG STREQUAL "CUDA") -cuda_kernel_component(SRC - NAME natten_hopper - SOURCES "csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_cm.h" -"csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_dtype.h" -"csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_head_dim.h" -"csrc/autogen/include/natten_autogen/cuda/hopper_fna/dispatch_tile_shape.h" -"csrc/autogen/include/natten_autogen/cuda/hopper_fna/interface.h" -"csrc/autogen/include/natten_autogen/cuda/hopper_fna/kernels.h" -"csrc/autogen/src/cuda/hopper_fna/source_0.cu" -"csrc/autogen/src/cuda/hopper_fna/source_1.cu" -"csrc/autogen/src/cuda/hopper_fna/source_2.cu" -"csrc/autogen/src/cuda/hopper_fna/source_3.cu" -"csrc/autogen/src/cuda/hopper_fna/source_4.cu" -"csrc/autogen/src/cuda/hopper_fna/source_5.cu" -"csrc/autogen/src/cuda/hopper_fna/source_6.cu" -"csrc/autogen/src/cuda/hopper_fna/source_7.cu" -"csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_cm.h" -"csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_dtype.h" -"csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_head_dim.h" -"csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/dispatch_tile_shape.h" -"csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/interface.h" -"csrc/autogen/include/natten_autogen/cuda/hopper_fna_bwd/kernels.h" -"csrc/autogen/src/cuda/hopper_fna_bwd/source_0.cu" -"csrc/autogen/src/cuda/hopper_fna_bwd/source_1.cu" -"csrc/autogen/src/cuda/hopper_fna_bwd/source_2.cu" -"csrc/autogen/src/cuda/hopper_fna_bwd/source_3.cu" -"csrc/autogen/include/natten_autogen/cuda/hopper_fmha/dispatch_head_dim.h" -"csrc/autogen/include/natten_autogen/cuda/hopper_fmha/dispatch_tile_size.h" -"csrc/autogen/include/natten_autogen/cuda/hopper_fmha/interface.h" -"csrc/autogen/include/natten_autogen/cuda/hopper_fmha/kernels.h" -"csrc/autogen/src/cuda/hopper_fmha/source_0.cu" -"csrc/autogen/src/cuda/hopper_fmha/source_1.cu" -"csrc/autogen/src/cuda/hopper_fmha/source_2.cu" -"csrc/autogen/src/cuda/hopper_fmha/source_3.cu" -"csrc/autogen/src/cuda/hopper_fmha/source_4.cu" -"csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/dispatch_head_dim.h" -"csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/dispatch_tile_size.h" -"csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/interface.h" -"csrc/autogen/include/natten_autogen/cuda/hopper_fmha_bwd/kernels.h" -"csrc/autogen/src/cuda/hopper_fmha_bwd/source_0.cu" -"csrc/autogen/src/cuda/hopper_fmha_bwd/source_1.cu" -"csrc/autogen/src/cuda/hopper_fmha_bwd/source_2.cu" -"csrc/autogen/src/cuda/hopper_fmha_bwd/source_3.cu" -"csrc/autogen/src/cuda/hopper_fmha_bwd/source_4.cu" - INCLUDES "${CMAKE_SOURCE_DIR}/csrc/include;${CMAKE_SOURCE_DIR}/csrc/autogen/include" CUDA_CAPABILITIES 9.0a CUDA_FLAGS "-O3;-std=c++17;--expt-relaxed-constexpr;--extended-lambda;--use_fast_math;--ftemplate-backtrace-limit=0;-Xfatbin=-compress-all;-DNDEBUG;-DNATTEN_WITH_CUTLASS;-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1;-DNATTEN_WITH_HOPPER_FNA=1" ) -endif() -# Include Metal shader compilation utilities if needed -if(GPU_LANG STREQUAL "METAL") - include(${CMAKE_CURRENT_LIST_DIR}/cmake/compile-metal.cmake) -endif() - -# Define the extension target with unified parameters -define_gpu_extension_target( - ${OPS_NAME} - ${OPS_NAME} - DESTINATION ${OPS_NAME} - LANGUAGE ${GPU_LANG} - SOURCES ${SRC} - COMPILE_FLAGS ${GPU_FLAGS} - USE_SABI 3 - WITH_SOABI) - -if(GPU_LANG STREQUAL "SYCL") - target_link_options(${OPS_NAME} PRIVATE ${sycl_link_flags}) - target_link_libraries(${OPS_NAME} PRIVATE dnnl) -endif() - -# Compile Metal shaders if any were found -if(GPU_LANG STREQUAL "METAL") - if(ALL_METAL_SOURCES) - compile_metal_shaders(${OPS_NAME} "${ALL_METAL_SOURCES}" "${METAL_INCLUDE_DIRS}") - endif() -endif() - - -# Add kernels_install target for huggingface/kernels library layout -add_kernels_install_target(${OPS_NAME} "natten" "${BUILD_VARIANT_NAME}" - DATA_EXTENSIONS "" - GPU_ARCHS "${ALL_GPU_ARCHS}") - -# Add local_install target for local development with get_local_kernel() -add_local_install_target(${OPS_NAME} "natten" "${BUILD_VARIANT_NAME}" - DATA_EXTENSIONS "" - GPU_ARCHS "${ALL_GPU_ARCHS}") diff --git a/natten/cmake/_ops.py.in b/natten/cmake/_ops.py.in deleted file mode 100644 index 736771e8..00000000 --- a/natten/cmake/_ops.py.in +++ /dev/null @@ -1,9 +0,0 @@ -import torch -from . import @OPS_NAME@ -ops = torch.ops.@OPS_NAME@ - -def add_op_namespace_prefix(op_name: str): - """ - Prefix op by namespace. - """ - return f"@OPS_NAME@::{op_name}" diff --git a/natten/cmake/add_build_metadata.py b/natten/cmake/add_build_metadata.py deleted file mode 100644 index 23a02492..00000000 --- a/natten/cmake/add_build_metadata.py +++ /dev/null @@ -1,56 +0,0 @@ -import argparse -import json -import sys - - -def main(): - parser = argparse.ArgumentParser( - description="Write a metadata JSON file with GPU architecture information, " - "reading from a source file and writing to a destination." - ) - parser.add_argument( - "input", - help="Path to the source metadata JSON file to read from.", - ) - parser.add_argument( - "destination", - help="Path to write the output metadata JSON file to.", - ) - - parser.add_argument( - "--archs", - help="Semicolon-separated list of GPU architectures/capabilities.", - ) - - args = parser.parse_args() - - archs = ( - sorted(set(a for a in args.archs.split(";") if a)) - if args.archs is not None - else None - ) - - try: - with open(args.input) as f: - data = json.load(f) - except FileNotFoundError: - print(f"Error: input metadata file not found: {args.input}", file=sys.stderr) - sys.exit(1) - except json.JSONDecodeError as e: - print(f"Error: failed to parse input metadata JSON: {e}", file=sys.stderr) - sys.exit(1) - - if archs is not None: - data["backend"]["archs"] = archs - - try: - with open(args.destination, "w") as f: - json.dump(data, f, indent=2) - f.write("\n") - except OSError as e: - print(f"Error: failed to write output metadata JSON: {e}", file=sys.stderr) - sys.exit(1) - - -if __name__ == "__main__": - main() diff --git a/natten/cmake/build-variants.cmake b/natten/cmake/build-variants.cmake deleted file mode 100644 index 31fb8526..00000000 --- a/natten/cmake/build-variants.cmake +++ /dev/null @@ -1,328 +0,0 @@ -# Generate a standardized build variant name following the pattern: -# torch-[cxx11-]-- -# or, when compiled against the Torch stable ABI: -# torch-stable-abi--- -# -# Arguments: -# OUT_BUILD_NAME - Output variable name -# TORCH_VERSION - PyTorch version (e.g., "2.7.1"); ignored when TORCH_STABLE_ABI is set -# COMPUTE_FRAMEWORK - One of: cuda, rocm, metal, xpu, cpu -# COMPUTE_VERSION - Version of compute framework (e.g., "12.4" for CUDA, "6.0" for ROCm) -# Optional for CPU-only builds (pass empty string or omit) -# Optional keyword arguments: -# TORCH_STABLE_ABI - Stable ABI version the extension was compiled against (e.g., "2.11"); -# when set, TORCH_VERSION is ignored and the prefix becomes -# torch-stable-abi (e.g., "2.11" -> "torch-stable-abi211") -# Example output: torch27-cxx11-cu124-x86_64-linux (Linux) -# torch27-cu124-x86_64-windows (Windows) -# torch27-metal-aarch64-darwin (macOS) -# torch-stable-abi211-cu124-x86_64-linux (Linux, stable ABI) -# -function(generate_build_name OUT_BUILD_NAME TORCH_VERSION COMPUTE_FRAMEWORK COMPUTE_VERSION) - cmake_parse_arguments(ARG "" "TORCH_STABLE_ABI" "" ${ARGN}) - - if(ARG_TORCH_STABLE_ABI) - # Flatten the stable ABI version (e.g., "2.11" -> "211") and ignore TORCH_VERSION - string(REPLACE "." ";" VERSION_LIST "${ARG_TORCH_STABLE_ABI}") - list(LENGTH VERSION_LIST VERSION_COMPONENTS) - if(VERSION_COMPONENTS LESS 2) - list(APPEND VERSION_LIST "0") - endif() - list(GET VERSION_LIST 0 MAJOR) - list(GET VERSION_LIST 1 MINOR) - set(TORCH_PREFIX "torch-stable-abi${MAJOR}${MINOR}") - else() - # Flatten version by removing dots and padding to 2 components - string(REPLACE "." ";" VERSION_LIST "${TORCH_VERSION}") - list(LENGTH VERSION_LIST VERSION_COMPONENTS) - - # Pad to at least 2 components - if(VERSION_COMPONENTS LESS 2) - list(APPEND VERSION_LIST "0") - endif() - - # Take first 2 components and join without dots - list(GET VERSION_LIST 0 MAJOR) - list(GET VERSION_LIST 1 MINOR) - set(TORCH_PREFIX "torch${MAJOR}${MINOR}") - endif() - - # Generate compute string - if(COMPUTE_FRAMEWORK STREQUAL "cuda") - # Flatten CUDA version (e.g., "12.4" -> "124") - string(REPLACE "." ";" COMPUTE_VERSION_LIST "${COMPUTE_VERSION}") - list(LENGTH COMPUTE_VERSION_LIST COMPUTE_COMPONENTS) - if(COMPUTE_COMPONENTS GREATER_EQUAL 2) - list(GET COMPUTE_VERSION_LIST 0 COMPUTE_MAJOR) - list(GET COMPUTE_VERSION_LIST 1 COMPUTE_MINOR) - set(COMPUTE_STRING "cu${COMPUTE_MAJOR}${COMPUTE_MINOR}") - else() - list(GET COMPUTE_VERSION_LIST 0 COMPUTE_MAJOR) - set(COMPUTE_STRING "cu${COMPUTE_MAJOR}0") - endif() - elseif(COMPUTE_FRAMEWORK STREQUAL "rocm") - # Flatten ROCm version (e.g., "6.0" -> "60") - string(REPLACE "." ";" COMPUTE_VERSION_LIST "${COMPUTE_VERSION}") - list(LENGTH COMPUTE_VERSION_LIST COMPUTE_COMPONENTS) - if(COMPUTE_COMPONENTS GREATER_EQUAL 2) - list(GET COMPUTE_VERSION_LIST 0 COMPUTE_MAJOR) - list(GET COMPUTE_VERSION_LIST 1 COMPUTE_MINOR) - set(COMPUTE_STRING "rocm${COMPUTE_MAJOR}${COMPUTE_MINOR}") - else() - list(GET COMPUTE_VERSION_LIST 0 COMPUTE_MAJOR) - set(COMPUTE_STRING "rocm${COMPUTE_MAJOR}0") - endif() - elseif(COMPUTE_FRAMEWORK STREQUAL "xpu") - # Flatten XPU version (e.g., "2025.2" -> "202552") - string(REPLACE "." ";" COMPUTE_VERSION_LIST "${COMPUTE_VERSION}") - list(LENGTH COMPUTE_VERSION_LIST COMPUTE_COMPONENTS) - if(COMPUTE_COMPONENTS GREATER_EQUAL 2) - list(GET COMPUTE_VERSION_LIST 0 COMPUTE_MAJOR) - list(GET COMPUTE_VERSION_LIST 1 COMPUTE_MINOR) - set(COMPUTE_STRING "xpu${COMPUTE_MAJOR}${COMPUTE_MINOR}") - else() - list(GET COMPUTE_VERSION_LIST 0 COMPUTE_MAJOR) - set(COMPUTE_STRING "xpu${COMPUTE_MAJOR}0") - endif() - elseif(COMPUTE_FRAMEWORK STREQUAL "metal") - set(COMPUTE_STRING "metal") - elseif(COMPUTE_FRAMEWORK STREQUAL "cpu") - set(COMPUTE_STRING "cpu") - else() - message(FATAL_ERROR "Unknown compute framework: ${COMPUTE_FRAMEWORK}") - endif() - - # Detect from target system (CMAKE_SYSTEM_* variables refer to target, not host) - # Normalize architecture name - if(CMAKE_SYSTEM_PROCESSOR MATCHES "^(x86_64|amd64|AMD64)$") - set(CPU_ARCH "x86_64") - elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "^(aarch64|arm64|ARM64)$") - set(CPU_ARCH "aarch64") - else() - message(FATAL_ERROR "Unsupported architecture: ${CMAKE_SYSTEM_PROCESSOR}") - endif() - - # Normalize OS name - if(CMAKE_SYSTEM_NAME STREQUAL "Windows") - set(OS_NAME "windows") - elseif(CMAKE_SYSTEM_NAME STREQUAL "Linux") - set(OS_NAME "linux") - elseif(CMAKE_SYSTEM_NAME STREQUAL "Darwin") - set(OS_NAME "darwin") - else() - message(WARNING "Unknown OS ${CMAKE_SYSTEM_NAME}, using as-is") - string(TOLOWER "${CMAKE_SYSTEM_NAME}" OS_NAME) - endif() - - set(ARCH_OS_STRING "${CPU_ARCH}-${OS_NAME}") - - # Assemble the final build name - # For non-stable-ABI Linux builds, include cxx11 ABI indicator for compatibility - if(NOT ARG_TORCH_STABLE_ABI AND ARCH_OS_STRING MATCHES "-linux$") - set(BUILD_NAME "${TORCH_PREFIX}-cxx11-${COMPUTE_STRING}-${ARCH_OS_STRING}") - else() - set(BUILD_NAME "${TORCH_PREFIX}-${COMPUTE_STRING}-${ARCH_OS_STRING}") - endif() - - set(${OUT_BUILD_NAME} "${BUILD_NAME}" PARENT_SCOPE) - message(STATUS "Generated build name: ${BUILD_NAME}") -endfunction() - -# -# Create a custom install target for the huggingface/kernels library layout. -# This installs the extension into a directory structure suitable for kernel hub discovery: -# / -# -# Arguments: -# TARGET_NAME - Name of the target to create the install rule for -# PACKAGE_NAME - Python package name (e.g., "activation") -# BUILD_VARIANT_NAME - Build variant name (e.g., "torch271-cxx11-cu124-x86_64-linux") -# INSTALL_PREFIX - Base installation directory (defaults to CMAKE_INSTALL_PREFIX) -# GPU_ARCHS - List of GPU architectures that were compiled -# (optional; when provided for CUDA/ROCm, metadata.json will include -# a "backend" key with the type and arch list) -# -function(add_kernels_install_target TARGET_NAME PACKAGE_NAME BUILD_VARIANT_NAME) - set(oneValueArgs INSTALL_PREFIX) - set(multiValueArgs DATA_EXTENSIONS GPU_ARCHS) - cmake_parse_arguments(ARG "" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) - - if(NOT ARG_INSTALL_PREFIX) - set(ARG_INSTALL_PREFIX "${CMAKE_INSTALL_PREFIX}") - endif() - - gpu_lang_to_backend(_BACKEND ${GPU_LANG}) - - # Always include 'py' extension for Python files - set(ALL_EXTENSIONS ${ARG_DATA_EXTENSIONS}) - list(APPEND ALL_EXTENSIONS "py") - - # Set the installation directory - set(KERNEL_INSTALL_DIR "${ARG_INSTALL_PREFIX}/${BUILD_VARIANT_NAME}") - - message(STATUS "Using PACKAGE_NAME: ${PACKAGE_NAME}") - - # Install the compiled extension using CMake's install() command - # This will be triggered by the standard INSTALL target - install(TARGETS ${TARGET_NAME} - LIBRARY DESTINATION "${KERNEL_INSTALL_DIR}" - RUNTIME DESTINATION "${KERNEL_INSTALL_DIR}" - COMPONENT ${TARGET_NAME}) - - # Install data files with specified extensions - foreach(ext IN LISTS ALL_EXTENSIONS) - file(GLOB_RECURSE DATA_FILES RELATIVE "${CMAKE_SOURCE_DIR}/torch-ext/${PACKAGE_NAME}" "${CMAKE_SOURCE_DIR}/torch-ext/${PACKAGE_NAME}/*.${ext}") - foreach(data_file IN LISTS DATA_FILES) - get_filename_component(data_file_dir "${data_file}" DIRECTORY) - install(FILES "${CMAKE_SOURCE_DIR}/torch-ext/${PACKAGE_NAME}/${data_file}" - DESTINATION "${KERNEL_INSTALL_DIR}/${data_file_dir}" - COMPONENT ${TARGET_NAME}) - endforeach() - endforeach() - - message(STATUS "GPU archs: ${ARG_GPU_ARCHS}") - - # Add the GPU archs to matadata.json when applicable. - if((GPU_LANG STREQUAL "CUDA" OR GPU_LANG STREQUAL "HIP") AND ARG_GPU_ARCHS) - list(JOIN ARG_GPU_ARCHS ";" _GPU_ARCHS_STR) - install(CODE " - file(MAKE_DIRECTORY \"${KERNEL_INSTALL_DIR}\") - execute_process( - COMMAND \"${Python3_EXECUTABLE}\" - \"${CMAKE_CURRENT_LIST_DIR}/cmake/add_build_metadata.py\" - \"${CMAKE_SOURCE_DIR}/metadata-${_BACKEND}.json\" - \"${KERNEL_INSTALL_DIR}/metadata.json\" - --archs \"${_GPU_ARCHS_STR}\" - RESULT_VARIABLE _METADATA_RESULT - ERROR_VARIABLE _METADATA_ERROR - ) - if(NOT _METADATA_RESULT EQUAL 0) - message(WARNING \"Failed to add GPU arch metadata: \${_METADATA_ERROR}\") - endif() - " COMPONENT ${TARGET_NAME}) - else() - install(CODE " - file(MAKE_DIRECTORY \"${KERNEL_INSTALL_DIR}\") - execute_process( - COMMAND \"${Python3_EXECUTABLE}\" - \"${CMAKE_CURRENT_LIST_DIR}/cmake/add_build_metadata.py\" - \"${CMAKE_SOURCE_DIR}/metadata-${_BACKEND}.json\" - \"${KERNEL_INSTALL_DIR}/metadata.json\" - RESULT_VARIABLE _METADATA_RESULT - ERROR_VARIABLE _METADATA_ERROR - ) - if(NOT _METADATA_RESULT EQUAL 0) - message(WARNING \"Failed to write metadata: \${_METADATA_ERROR}\") - endif() - " COMPONENT ${TARGET_NAME}) - endif() - - # Compatibility with older kernels and direct Python imports. - install(FILES ${CMAKE_SOURCE_DIR}/compat.py - DESTINATION "${KERNEL_INSTALL_DIR}/${PACKAGE_NAME}" - RENAME "__init__.py" - COMPONENT ${TARGET_NAME}) - - message(STATUS "Added install rules for ${TARGET_NAME} -> ${BUILD_VARIANT_NAME}") -endfunction() - -# -# Add install rules for local development with huggingface/kernels. -# This installs the extension into the layout expected by get_local_kernel(): -# ${CMAKE_SOURCE_DIR}/build// -# -# This allows developers to use get_local_kernel() from the kernels library to load -# locally built kernels without needing to publish to the hub. -# -# This uses the standard CMake install() command, so it works with the default -# "install" target that is always available. -# -# Arguments: -# TARGET_NAME - Name of the target to create the install rule for -# PACKAGE_NAME - Python package name (e.g., "activation") -# BUILD_VARIANT_NAME - Build variant name (e.g., "torch271-cxx11-cu124-x86_64-linux") -# GPU_ARCHS - List of GPU architectures that were compiled -# (optional; when provided for CUDA/ROCm, metadata.json will include -# a "backend" key with the type and arch list) -# -function(add_local_install_target TARGET_NAME PACKAGE_NAME BUILD_VARIANT_NAME) - set(oneValueArgs) - set(multiValueArgs DATA_EXTENSIONS GPU_ARCHS) - cmake_parse_arguments(ARG "" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) - - # Always include 'py' extension for Python files - set(ALL_EXTENSIONS ${ARG_DATA_EXTENSIONS}) - list(APPEND ALL_EXTENSIONS "py") - - # Define your local, folder based, installation directory - set(LOCAL_INSTALL_DIR "${CMAKE_SOURCE_DIR}/build/${BUILD_VARIANT_NAME}") - # Variant directory is where metadata.json should go (for kernels upload discovery) - set(VARIANT_DIR "${CMAKE_SOURCE_DIR}/build/${BUILD_VARIANT_NAME}") - - # Create a custom target for local installation - add_custom_target(local_install - COMMENT "Installing files to local directory..." - ) - - gpu_lang_to_backend(_BACKEND ${GPU_LANG}) - - # Copy data files with specified extensions - foreach(ext IN LISTS ALL_EXTENSIONS) - file(GLOB_RECURSE DATA_FILES RELATIVE "${CMAKE_SOURCE_DIR}/torch-ext/${PACKAGE_NAME}" "${CMAKE_SOURCE_DIR}/torch-ext/${PACKAGE_NAME}/*.${ext}") - foreach(data_file IN LISTS DATA_FILES) - get_filename_component(data_file_dir "${data_file}" DIRECTORY) - add_custom_command(TARGET local_install POST_BUILD - COMMAND ${CMAKE_COMMAND} -E make_directory - ${LOCAL_INSTALL_DIR}/${data_file_dir} - COMMAND ${CMAKE_COMMAND} -E copy_if_different - ${CMAKE_SOURCE_DIR}/torch-ext/${PACKAGE_NAME}/${data_file} - ${LOCAL_INSTALL_DIR}/${data_file_dir}/ - COMMENT "Copying ${data_file} to ${LOCAL_INSTALL_DIR}/${data_file_dir}" - ) - endforeach() - endforeach() - - - # Add the GPU archs to matadata.json when applicable. - if((GPU_LANG STREQUAL "CUDA" OR GPU_LANG STREQUAL "HIP") AND ARG_GPU_ARCHS) - list(JOIN ARG_GPU_ARCHS ";" _GPU_ARCHS_STR) - add_custom_command(TARGET local_install POST_BUILD - COMMAND ${CMAKE_COMMAND} -E make_directory ${VARIANT_DIR} - COMMAND ${Python3_EXECUTABLE} - ${CMAKE_CURRENT_LIST_DIR}/cmake/add_build_metadata.py - ${CMAKE_SOURCE_DIR}/metadata-${_BACKEND}.json - ${VARIANT_DIR}/metadata.json - --archs "${_GPU_ARCHS_STR}" - COMMENT "Writing metadata.json with GPU arch info to ${VARIANT_DIR}" - ) - else() - add_custom_command(TARGET local_install POST_BUILD - COMMAND ${CMAKE_COMMAND} -E make_directory ${VARIANT_DIR} - COMMAND ${Python3_EXECUTABLE} - ${CMAKE_CURRENT_LIST_DIR}/cmake/add_build_metadata.py - ${CMAKE_SOURCE_DIR}/metadata-${_BACKEND}.json - ${VARIANT_DIR}/metadata.json - COMMENT "Writing metadata.json to ${VARIANT_DIR}" - ) - endif() - - add_custom_command(TARGET local_install POST_BUILD - # Copy the shared library - COMMAND ${CMAKE_COMMAND} -E copy_if_different - $ - ${LOCAL_INSTALL_DIR}/ - - # Compatibility with older kernels and direct Python imports. - COMMAND ${CMAKE_COMMAND} -E copy_if_different - ${CMAKE_SOURCE_DIR}/compat.py - ${VARIANT_DIR}/${PACKAGE_NAME}/__init__.py - - COMMENT "Copying shared library and Python files to ${LOCAL_INSTALL_DIR}" - COMMAND_EXPAND_LISTS - ) - - # Create both directories: variant dir for metadata.json, package dir for binaries - file(MAKE_DIRECTORY ${VARIANT_DIR}) - file(MAKE_DIRECTORY ${LOCAL_INSTALL_DIR}) - message(STATUS "Added install rules for ${TARGET_NAME} -> build/${BUILD_VARIANT_NAME}") -endfunction() diff --git a/natten/cmake/compile-metal.cmake b/natten/cmake/compile-metal.cmake deleted file mode 100644 index 50d44a2d..00000000 --- a/natten/cmake/compile-metal.cmake +++ /dev/null @@ -1,104 +0,0 @@ -# Metal shader compilation function -function(compile_metal_shaders TARGET_NAME METAL_SOURCES EXTRA_INCLUDE_DIRS) - if(NOT DEFINED METAL_TOOLCHAIN) - execute_process( - COMMAND "xcodebuild" "-showComponent" "MetalToolchain" - OUTPUT_VARIABLE FIND_METAL_OUT - RESULT_VARIABLE FIND_METAL_ERROR_CODE - ERROR_VARIABLE FIND_METAL_STDERR - OUTPUT_STRIP_TRAILING_WHITESPACE) - - if(NOT FIND_METAL_ERROR_CODE EQUAL 0) - message(FATAL_ERROR "${ERR_MSG}: ${FIND_METAL_STDERR}") - endif() - - # Extract the Toolchain Search Path value and append Metal.xctoolchain - string(REGEX MATCH "Toolchain Search Path: ([^\n]+)" MATCH_RESULT "${FIND_METAL_OUT}") - set(METAL_TOOLCHAIN "${CMAKE_MATCH_1}/Metal.xctoolchain") - endif() - - # Set Metal compiler flags - set(METAL_FLAGS "-std=metal4.0" "-O2") - - # Output directory for compiled metallib - set(METALLIB_OUTPUT_DIR "${CMAKE_BINARY_DIR}/metallib") - file(MAKE_DIRECTORY ${METALLIB_OUTPUT_DIR}) - - foreach(INC ${EXTRA_INCLUDE_DIRS}) - list(APPEND METAL_FLAGS "-I${INC}") - endforeach() - - # Separate .metal files from .h files and compile .metal files to .air - set(AIR_FILES) - set(METAL_FILES) - set(HEADER_FILES) - - foreach(SOURCE_FILE ${METAL_SOURCES}) - if(SOURCE_FILE MATCHES "\\.metal$") - list(APPEND METAL_FILES ${SOURCE_FILE}) - elseif(SOURCE_FILE MATCHES "\\.h$") - list(APPEND HEADER_FILES ${SOURCE_FILE}) - endif() - endforeach() - - foreach(METAL_FILE ${METAL_FILES}) - get_filename_component(METAL_NAME ${METAL_FILE} NAME_WE) - set(AIR_FILE "${CMAKE_BINARY_DIR}/${METAL_NAME}.air") - - # Include header files as dependencies - set(ALL_DEPENDENCIES ${CMAKE_CURRENT_SOURCE_DIR}/${METAL_FILE}) - foreach(HEADER_FILE ${HEADER_FILES}) - list(APPEND ALL_DEPENDENCIES ${CMAKE_CURRENT_SOURCE_DIR}/${HEADER_FILE}) - endforeach() - - add_custom_command( - OUTPUT ${AIR_FILE} - COMMAND "${METAL_TOOLCHAIN}/usr/bin/metal" ${METAL_FLAGS} - -c ${CMAKE_CURRENT_SOURCE_DIR}/${METAL_FILE} - -o ${AIR_FILE} - DEPENDS ${ALL_DEPENDENCIES} - COMMENT "Compiling Metal shader ${METAL_FILE} to ${AIR_FILE}" - VERBATIM - ) - - list(APPEND AIR_FILES ${AIR_FILE}) - endforeach() - - # Link all .air files into a single .metallib - set(METALLIB_FILE "${METALLIB_OUTPUT_DIR}/${TARGET_NAME}.metallib") - add_custom_command( - OUTPUT ${METALLIB_FILE} - COMMAND "${METAL_TOOLCHAIN}/usr/bin/metallib" ${AIR_FILES} - -o ${METALLIB_FILE} - DEPENDS ${AIR_FILES} - COMMENT "Linking Metal library ${METALLIB_FILE}" - VERBATIM - ) - - # Generate C++ header with embedded metallib data - set(METALLIB_HEADER "${CMAKE_BINARY_DIR}/${TARGET_NAME}_metallib.h") - set(METALLIB_TO_HEADER_SCRIPT "${CMAKE_CURRENT_SOURCE_DIR}/cmake/metallib_to_header.py") - - add_custom_command( - OUTPUT ${METALLIB_HEADER} - COMMAND ${Python3_EXECUTABLE} ${METALLIB_TO_HEADER_SCRIPT} ${METALLIB_FILE} ${METALLIB_HEADER} ${TARGET_NAME} - DEPENDS ${METALLIB_FILE} ${METALLIB_TO_HEADER_SCRIPT} - COMMENT "Generating embedded Metal library header ${METALLIB_HEADER}" - VERBATIM - ) - - # Create a custom target for the metallib - add_custom_target(${TARGET_NAME}_metallib ALL DEPENDS ${METALLIB_FILE} ${METALLIB_HEADER}) - - # Add dependency to main target - add_dependencies(${TARGET_NAME} ${TARGET_NAME}_metallib) - - # Add the generated header to include directories - target_include_directories(${TARGET_NAME} PRIVATE ${CMAKE_BINARY_DIR}) - - # Pass the metallib header and namespace as compile definitions - target_compile_definitions(${TARGET_NAME} PRIVATE - EMBEDDED_METALLIB_HEADER="${TARGET_NAME}_metallib.h" - EMBEDDED_METALLIB_NAMESPACE=${TARGET_NAME}_metal - ) -endfunction() diff --git a/natten/cmake/get_gpu_lang.cmake b/natten/cmake/get_gpu_lang.cmake deleted file mode 100644 index 277ac7b0..00000000 --- a/natten/cmake/get_gpu_lang.cmake +++ /dev/null @@ -1,9 +0,0 @@ -# -# Get the GPU language from Torch. -# -function(get_gpu_lang OUT) - run_python_script(PYTHON_OUT - "${CMAKE_CURRENT_SOURCE_DIR}/cmake/get_gpu_lang.py" - "Cannot detect GPU language") - set(${OUT} ${PYTHON_OUT} PARENT_SCOPE) -endfunction() diff --git a/natten/cmake/get_gpu_lang.py b/natten/cmake/get_gpu_lang.py deleted file mode 100644 index 1eedff7e..00000000 --- a/natten/cmake/get_gpu_lang.py +++ /dev/null @@ -1,20 +0,0 @@ -#!/usr/bin/env python3 - -import sys - -try: - import torch -except ImportError: - print("Torch is required for configuring a kernel build.", file=sys.stderr) - sys.exit(1) - -if torch.version.cuda is not None: - print("CUDA") -elif torch.version.hip is not None: - print("HIP") -elif torch.backends.mps.is_available(): - print("METAL") -elif hasattr(torch.version, "xpu") and torch.version.xpu is not None: - print("SYCL") -else: - print("CPU") diff --git a/natten/cmake/hipify.py b/natten/cmake/hipify.py deleted file mode 100644 index a1539c02..00000000 --- a/natten/cmake/hipify.py +++ /dev/null @@ -1,76 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-License-Identifier: Apache-2.0 - -# From vLLM: https://github.com/vllm-project/vllm/blob/main/cmake/hipify.py - -# -# A command line tool for running pytorch's hipify preprocessor on CUDA -# source files. -# -# See https://github.com/ROCm/hipify_torch -# and /utils/hipify/hipify_python.py -# - -import argparse -import os -import shutil - -from torch.utils.hipify.hipify_python import hipify - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - - # Project directory where all the source + include files live. - parser.add_argument( - "-p", - "--project_dir", - help="The project directory.", - ) - - # Directory where hipified files are written. - parser.add_argument( - "-o", - "--output_dir", - help="The output directory.", - ) - - # Source files to convert. - parser.add_argument("sources", - help="Source files to hipify.", - nargs="*", - default=[]) - - args = parser.parse_args() - - # Limit include scope to project_dir only - includes = [os.path.join(args.project_dir, '*')] - - # Get absolute path for all source files. - extra_files = [os.path.abspath(s) for s in args.sources] - - # Copy sources from project directory to output directory. - # The directory might already exist to hold object files so we ignore that. - shutil.copytree(args.project_dir, args.output_dir, dirs_exist_ok=True) - - hipify_result = hipify(project_directory=args.project_dir, - output_directory=args.output_dir, - header_include_dirs=[], - includes=includes, - extra_files=extra_files, - show_detailed=True, - is_pytorch_extension=True, - hipify_extra_files_only=True) - - hipified_sources = [] - for source in args.sources: - s_abs = os.path.abspath(source) - hipified_s_abs = (hipify_result[s_abs].hipified_path if - (s_abs in hipify_result - and hipify_result[s_abs].hipified_path is not None) - else s_abs) - hipified_sources.append(hipified_s_abs) - - assert (len(hipified_sources) == len(args.sources)) - - # Print hipified source files. - print("\n".join(hipified_sources)) diff --git a/natten/cmake/kernel.cmake b/natten/cmake/kernel.cmake deleted file mode 100644 index c41aaa65..00000000 --- a/natten/cmake/kernel.cmake +++ /dev/null @@ -1,317 +0,0 @@ -function(accumulate_gpu_archs OUT_ACC ACC EXTRA_ARCHS) - list(APPEND ACC ${EXTRA_ARCHS}) - list(REMOVE_DUPLICATES ACC) - list(SORT ACC) - set(${OUT_ACC} ${ACC} PARENT_SCOPE) -endfunction() - -function(cuda_kernel_component SRC_VAR) - set(oneValueArgs CUDA_MINVER NAME) - set(multiValueArgs SOURCES INCLUDES CUDA_CAPABILITIES CUDA_FLAGS CXX_FLAGS) - cmake_parse_arguments(KERNEL "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) - - if(NOT KERNEL_SOURCES) - message(FATAL_ERROR "cuda_kernel_component: SOURCES argument is required") - endif() - - # Bail out if this component is not supported by the CUDA version. - if(KERNEL_CUDA_MINVER) - if(CUDA_VERSION VERSION_LESS ${KERNEL_CUDA_MINVER}) - return() - endif() - endif() - - set(_KERNEL_SRC ${KERNEL_SOURCES}) - - if(KERNEL_INCLUDES) - # TODO: check if CLion support this: - # https://youtrack.jetbrains.com/issue/CPP-16510/CLion-does-not-handle-per-file-include-directories - set_source_files_properties( - ${_KERNEL_SRC} - PROPERTIES INCLUDE_DIRECTORIES "${KERNEL_INCLUDES}") - endif() - - # Determine CUDA architectures - if(KERNEL_CUDA_CAPABILITIES) - cuda_archs_loose_intersection(_KERNEL_ARCHS "${KERNEL_CUDA_CAPABILITIES}" "${CUDA_ARCHS}") - if(NOT _KERNEL_ARCHS) - message(FATAL_ERROR "CUDA kernel: ${KERNEL_NAME}, empty set of capabilities after intersection (kernel: ${KERNEL_CUDA_CAPABILITIES}, supported: ${CUDA_ARCHS})") - endif() - else() - set(_KERNEL_ARCHS "${CUDA_KERNEL_ARCHS}") - endif() - message(STATUS "CUDA kernel: ${KERNEL_NAME}, capabilities: ${_KERNEL_ARCHS}") - set_gencode_flags_for_srcs(SRCS "${_KERNEL_SRC}" CUDA_ARCHS "${_KERNEL_ARCHS}") - - accumulate_gpu_archs(_ALL_GPU_ARCHS "${ALL_GPU_ARCHS}" "${_KERNEL_ARCHS}") - set(ALL_GPU_ARCHS ${_ALL_GPU_ARCHS} PARENT_SCOPE) - - # Apply CUDA-specific compile flags - if(KERNEL_CUDA_FLAGS) - set(_CUDA_FLAGS "${KERNEL_CUDA_FLAGS}") - # -static-global-template-stub is not supported on CUDA < 12.8. Remove this - # once we don't support CUDA 12.6 anymore. - if(CUDA_VERSION VERSION_LESS 12.8) - string(REGEX REPLACE "-static-global-template-stub=(true|false)" "" _CUDA_FLAGS "${_CUDA_FLAGS}") - endif() - - foreach(_SRC ${_KERNEL_SRC}) - if(_SRC MATCHES ".*\\.cu$") - set_property( - SOURCE ${_SRC} - APPEND PROPERTY - COMPILE_OPTIONS "$<$:${_CUDA_FLAGS}>" - ) - endif() - endforeach() - endif() - - # Apply CXX-specific compile flags - if(KERNEL_CXX_FLAGS) - foreach(_SRC ${_KERNEL_SRC}) - set_property( - SOURCE ${_SRC} - APPEND PROPERTY - COMPILE_OPTIONS "$<$:${KERNEL_CXX_FLAGS}>" - ) - endforeach() - endif() - - set(_TMP_SRC ${${SRC_VAR}}) - list(APPEND _TMP_SRC ${_KERNEL_SRC}) - set(${SRC_VAR} ${_TMP_SRC} PARENT_SCOPE) -endfunction() - -function(hip_kernel_component SRC_VAR) - set(options SUPPORTS_HIPIFY) - set(oneValueArgs CUDA_MINVER NAME) - set(multiValueArgs SOURCES INCLUDES CXX_FLAGS HIP_FLAGS ROCM_ARCHS) - cmake_parse_arguments(KERNEL "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) - - if(NOT KERNEL_SOURCES) - message(FATAL_ERROR "hip_kernel_component: SOURCES argument is required") - endif() - - set(_KERNEL_SRC ${KERNEL_SOURCES}) - - if(KERNEL_INCLUDES) - # TODO: check if CLion support this: - # https://youtrack.jetbrains.com/issue/CPP-16510/CLion-does-not-handle-per-file-include-directories - set_source_files_properties( - ${_KERNEL_SRC} - PROPERTIES INCLUDE_DIRECTORIES "${KERNEL_INCLUDES}") - endif() - - # Apply HIP-specific compile flags - if(KERNEL_HIP_FLAGS) - foreach(_SRC ${_KERNEL_SRC}) - if(_SRC MATCHES ".*\\.(cu|hip)$") - set_property( - SOURCE ${_SRC} - APPEND PROPERTY - COMPILE_OPTIONS "$<$:${KERNEL_HIP_FLAGS}>" - ) - endif() - endforeach() - endif() - - # Determine ROCm architectures - if(KERNEL_ROCM_ARCHS) - hip_archs_loose_intersection(_KERNEL_ARCHS "${KERNEL_ROCM_ARCHS}" "${ROCM_ARCHS}") - if(NOT _KERNEL_ARCHS) - message(FATAL_ERROR "ROCm kernel: ${KERNEL_NAME}, empty set of architectures after intersection (kernel: ${KERNEL_ROCM_ARCHS}, supported: ${ROCM_ARCHS})") - endif() - else() - set(_KERNEL_ARCHS "${ROCM_ARCHS}") - endif() - message(STATUS "ROCm kernel: ${KERNEL_NAME}, archs: ${_KERNEL_ARCHS}") - - accumulate_gpu_archs(_ALL_GPU_ARCHS "${ALL_GPU_ARCHS}" "${_KERNEL_ARCHS}") - set(ALL_GPU_ARCHS ${_ALL_GPU_ARCHS} PARENT_SCOPE) - - foreach(_SRC ${_KERNEL_SRC}) - if(_SRC MATCHES ".*\\.(cu|hip)$") - foreach(_ARCH ${_KERNEL_ARCHS}) - set_property( - SOURCE ${_SRC} - APPEND PROPERTY - COMPILE_OPTIONS "$<$:--offload-arch=${_ARCH}>" - ) - endforeach() - endif() - endforeach() - - set(_TMP_SRC ${${SRC_VAR}}) - list(APPEND _TMP_SRC ${_KERNEL_SRC}) - set(${SRC_VAR} ${_TMP_SRC} PARENT_SCOPE) -endfunction() - - -function(xpu_kernel_component SRC_VAR) - set(options) - set(oneValueArgs) - set(multiValueArgs SOURCES INCLUDES CXX_FLAGS SYCL_FLAGS) - cmake_parse_arguments(KERNEL "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) - - if(NOT KERNEL_SOURCES) - message(FATAL_ERROR "xpu_kernel_component: SOURCES argument is required") - endif() - - set(_KERNEL_SRC ${KERNEL_SOURCES}) - - # Handle per-file include directories if specified - if(KERNEL_INCLUDES) - # TODO: check if CLion support this: - # https://youtrack.jetbrains.com/issue/CPP-16510/CLion-does-not-handle-per-file-include-directories - set_source_files_properties( - ${_KERNEL_SRC} - PROPERTIES INCLUDE_DIRECTORIES "${KERNEL_INCLUDES}") - endif() - - # Apply CXX-specific compile flags - if(KERNEL_CXX_FLAGS) - foreach(_SRC ${_KERNEL_SRC}) - set_property( - SOURCE ${_SRC} - APPEND PROPERTY - COMPILE_OPTIONS "$<$:${KERNEL_CXX_FLAGS}>" - ) - endforeach() - endif() - - # Add SYCL-specific compilation flags for XPU sources - if(KERNEL_SYCL_FLAGS) - # Use kernel-specific SYCL flags plus SYCL flags from the parent scope. - set(_SYCL_FLAGS ${sycl_flags}) - list(APPEND _SYCL_FLAGS ${KERNEL_SYCL_FLAGS}) - foreach(_SRC ${_KERNEL_SRC}) - if(_SRC MATCHES ".*\\.(cpp|cxx|cc)$") - set_property( - SOURCE ${_SRC} - APPEND PROPERTY - COMPILE_OPTIONS "$<$:${_SYCL_FLAGS}>" - ) - endif() - endforeach() - else() - # Use default SYCL flags (from parent scope variable sycl_flags) - foreach(_SRC ${_KERNEL_SRC}) - if(_SRC MATCHES ".*\\.(cpp|cxx|cc)$") - set_property( - SOURCE ${_SRC} - APPEND PROPERTY - COMPILE_OPTIONS "$<$:${sycl_flags}>" - ) - endif() - endforeach() - endif() - - # Append to parent scope SRC variable - set(_TMP_SRC ${${SRC_VAR}}) - list(APPEND _TMP_SRC ${_KERNEL_SRC}) - set(${SRC_VAR} ${_TMP_SRC} PARENT_SCOPE) -endfunction() - -function(cpu_kernel_component SRC_VAR) - set(options) - set(oneValueArgs) - set(multiValueArgs SOURCES INCLUDES CXX_FLAGS) - cmake_parse_arguments(KERNEL "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) - - if(NOT KERNEL_SOURCES) - message(FATAL_ERROR "cpu_kernel_component: SOURCES argument is required") - endif() - - set(_KERNEL_SRC ${KERNEL_SOURCES}) - - # Handle per-file include directories if specified - if(KERNEL_INCLUDES) - # TODO: check if CLion support this: - # https://youtrack.jetbrains.com/issue/CPP-16510/CLion-does-not-handle-per-file-include-directories - set_source_files_properties( - ${_KERNEL_SRC} - PROPERTIES INCLUDE_DIRECTORIES "${KERNEL_INCLUDES}") - endif() - - # Apply CXX-specific compile flags - if(KERNEL_CXX_FLAGS) - foreach(_SRC ${_KERNEL_SRC}) - set_property( - SOURCE ${_SRC} - APPEND PROPERTY - COMPILE_OPTIONS "$<$:${KERNEL_CXX_FLAGS}>" - ) - endforeach() - endif() - - # Append to parent scope SRC variable - set(_TMP_SRC ${${SRC_VAR}}) - list(APPEND _TMP_SRC ${_KERNEL_SRC}) - set(${SRC_VAR} ${_TMP_SRC} PARENT_SCOPE) -endfunction() - -function(metal_kernel_component SRC_VAR) - set(options) - set(oneValueArgs) - set(multiValueArgs SOURCES INCLUDES CXX_FLAGS) - cmake_parse_arguments(KERNEL "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) - - if(NOT KERNEL_SOURCES) - message(FATAL_ERROR "metal_kernel_component: SOURCES argument is required") - endif() - - set(_KERNEL_SRC ${KERNEL_SOURCES}) - - # Separate Metal shader files from other sources - set(_METAL_SRC) - set(_CPP_SRC) - - foreach(_SRC_FILE IN LISTS _KERNEL_SRC) - if(_SRC_FILE MATCHES "\\.(metal|h)$") - list(APPEND _METAL_SRC ${_SRC_FILE}) - else() - list(APPEND _CPP_SRC ${_SRC_FILE}) - endif() - endforeach() - - # Handle per-file include directories if specified (for C++ sources only) - if(KERNEL_INCLUDES AND _CPP_SRC) - # TODO: check if CLion support this: - # https://youtrack.jetbrains.com/issue/CPP-16510/CLion-does-not-handle-per-file-include-directories - set_source_files_properties( - ${_CPP_SRC} - PROPERTIES INCLUDE_DIRECTORIES "${KERNEL_INCLUDES}") - endif() - - # Apply CXX-specific compile flags - if(KERNEL_CXX_FLAGS AND _CPP_SRC) - foreach(_SRC ${_CPP_SRC}) - set_property( - SOURCE ${_SRC} - APPEND PROPERTY - COMPILE_OPTIONS "$<$:${KERNEL_CXX_FLAGS}>" - ) - endforeach() - endif() - - # Add C++ sources to main source list - if(_CPP_SRC) - set(_TMP_SRC ${${SRC_VAR}}) - list(APPEND _TMP_SRC ${_CPP_SRC}) - set(${SRC_VAR} ${_TMP_SRC} PARENT_SCOPE) - endif() - - # Keep track of Metal sources for later compilation - if(_METAL_SRC) - set(_TMP_METAL ${ALL_METAL_SOURCES}) - list(APPEND _TMP_METAL ${_METAL_SRC}) - set(ALL_METAL_SOURCES ${_TMP_METAL} PARENT_SCOPE) - endif() - - # Keep the includes directory for the Metal sources - if(KERNEL_INCLUDES AND _METAL_SRC) - set(_TMP_METAL_INCLUDES ${METAL_INCLUDE_DIRS}) - list(APPEND _TMP_METAL_INCLUDES ${KERNEL_INCLUDES}) - set(METAL_INCLUDE_DIRS ${_TMP_METAL_INCLUDES} PARENT_SCOPE) - endif() -endfunction() diff --git a/natten/cmake/metallib_to_header.py b/natten/cmake/metallib_to_header.py deleted file mode 100644 index 82bd252b..00000000 --- a/natten/cmake/metallib_to_header.py +++ /dev/null @@ -1,73 +0,0 @@ -#!/usr/bin/env python3 -import sys -import os - -def convert_metallib_to_header(metallib_path: str, header_path: str, target_name: str) -> None: - """Convert a metallib binary file to a C++ header with embedded data.""" - - # Read the metallib binary data - with open(metallib_path, 'rb') as f: - data: bytes = f.read() - - # Generate the header content - header_content: str = """// Auto-generated file containing embedded Metal library -#pragma once -#include -#include - -namespace """ + target_name + """_metal { - static const unsigned char metallib_data[] = { -""" - - # Convert binary data to C array format - bytes_per_line: int = 16 - for i in range(0, len(data), bytes_per_line): - chunk: bytes = data[i:i + bytes_per_line] - hex_values: str = ', '.join('0x{:02x}'.format(b) for b in chunk) - header_content += " " + hex_values + "," - if i + bytes_per_line < len(data): - header_content += "\n" - - header_content += """ - }; - static const size_t metallib_data_len = """ + str(len(data)) + """; - - // Convenience function to create Metal library from embedded data - inline id createLibrary(id device, NSError** error = nullptr) { - dispatch_data_t libraryData = dispatch_data_create( - metallib_data, - metallib_data_len, - dispatch_get_main_queue(), - ^{ /* No cleanup needed for static data */ }); - - NSError* localError = nil; - id library = [device newLibraryWithData:libraryData error:&localError]; - - if (error) { - *error = localError; - } - - return library; - } -} // namespace """ + target_name + """_metal -""" - - # Write the header file - dir_path: str = os.path.dirname(header_path) - if dir_path: - os.makedirs(dir_path, exist_ok=True) - with open(header_path, 'w') as f: - f.write(header_content) - - print("Generated {} ({} bytes)".format(header_path, len(data))) - -if __name__ == "__main__": - if len(sys.argv) != 4: - print("Usage: metallib_to_header.py ") - sys.exit(1) - - metallib_path: str = sys.argv[1] - header_path: str = sys.argv[2] - target_name: str = sys.argv[3] - - convert_metallib_to_header(metallib_path, header_path, target_name) \ No newline at end of file diff --git a/natten/cmake/utils.cmake b/natten/cmake/utils.cmake deleted file mode 100644 index 4c19c65c..00000000 --- a/natten/cmake/utils.cmake +++ /dev/null @@ -1,611 +0,0 @@ -# Vendored from vLLM: -# -# https://github.com/vllm-project/vllm/blob/main/cmake/utils.cmake -# -# Attempt to find the python package that uses the same python executable as -# `EXECUTABLE` and is one of the `SUPPORTED_VERSIONS`. -# -macro (find_python_from_executable EXECUTABLE SUPPORTED_VERSIONS) - file(REAL_PATH ${EXECUTABLE} EXECUTABLE) - set(Python3_EXECUTABLE ${EXECUTABLE}) - find_package(Python3 COMPONENTS Interpreter Development.Module Development.SABIModule) - if (NOT Python3_FOUND) - message(FATAL_ERROR "Unable to find python matching: ${EXECUTABLE}.") - endif() - set(_VER "${Python3_VERSION_MAJOR}.${Python3_VERSION_MINOR}") - set(_SUPPORTED_VERSIONS_LIST ${SUPPORTED_VERSIONS} ${ARGN}) - if (NOT _VER IN_LIST _SUPPORTED_VERSIONS_LIST) - message(FATAL_ERROR - "Python version (${_VER}) is not one of the supported versions: " - "${_SUPPORTED_VERSIONS_LIST}.") - endif() - message(STATUS "Found python matching: ${EXECUTABLE}.") -endmacro() - -# -# Run `EXPR` in python. The standard output of python is stored in `OUT` and -# has trailing whitespace stripped. If an error is encountered when running -# python, a fatal message `ERR_MSG` is issued. -# -function (run_python OUT EXPR ERR_MSG) - if(Python3_EXECUTABLE) - set(_PYTHON_EXECUTABLE "${Python3_EXECUTABLE}") - elseif(Python_EXECUTABLE) - set(_PYTHON_EXECUTABLE "${Python_EXECUTABLE}") - else() - message(FATAL_ERROR "No Python executable found. Set Python3_EXECUTABLE or Python_EXECUTABLE.") - endif() - execute_process( - COMMAND - "${_PYTHON_EXECUTABLE}" "-c" "${EXPR}" - OUTPUT_VARIABLE PYTHON_OUT - RESULT_VARIABLE PYTHON_ERROR_CODE - ERROR_VARIABLE PYTHON_STDERR - OUTPUT_STRIP_TRAILING_WHITESPACE) - - if(NOT PYTHON_ERROR_CODE EQUAL 0) - message(FATAL_ERROR "${ERR_MSG}: ${PYTHON_STDERR}") - endif() - set(${OUT} ${PYTHON_OUT} PARENT_SCOPE) -endfunction() - -# -# Run `SCRIPT_PATH` in Python. The standard output of Python is stored in -# `OUT` and has trailing whitespace stripped. If the script exits with a -# non-zero code, a fatal message `ERR_MSG` is issued. -# -function (run_python_script OUT SCRIPT_PATH ERR_MSG) - if(Python3_EXECUTABLE) - set(_PYTHON_EXECUTABLE "${Python3_EXECUTABLE}") - elseif(Python_EXECUTABLE) - set(_PYTHON_EXECUTABLE "${Python_EXECUTABLE}") - else() - message(FATAL_ERROR "No Python executable found. Set Python3_EXECUTABLE or Python_EXECUTABLE.") - endif() - execute_process( - COMMAND - "${_PYTHON_EXECUTABLE}" "${SCRIPT_PATH}" - OUTPUT_VARIABLE PYTHON_OUT - RESULT_VARIABLE PYTHON_ERROR_CODE - ERROR_VARIABLE PYTHON_STDERR - OUTPUT_STRIP_TRAILING_WHITESPACE) - - if(NOT PYTHON_ERROR_CODE EQUAL 0) - message(FATAL_ERROR "${ERR_MSG}: ${PYTHON_STDERR}") - endif() - set(${OUT} ${PYTHON_OUT} PARENT_SCOPE) -endfunction() - -# -# Run `EXPR` in python. The standard output of python is stored in `OUT` and -# has trailing whitespace stripped. If an error is encountered when running -# python, `SUCCESS` is set to FALSE. If successful, `SUCCESS` is set to TRUE. -# -function (try_run_python OUT SUCCESS EXPR) - if(Python3_EXECUTABLE) - set(_PYTHON_EXECUTABLE "${Python3_EXECUTABLE}") - elseif(Python_EXECUTABLE) - set(_PYTHON_EXECUTABLE "${Python_EXECUTABLE}") - else() - message(FATAL_ERROR "No Python executable found. Set Python3_EXECUTABLE or Python_EXECUTABLE.") - endif() - execute_process( - COMMAND - "${_PYTHON_EXECUTABLE}" "-c" "${EXPR}" - OUTPUT_VARIABLE PYTHON_OUT - RESULT_VARIABLE PYTHON_ERROR_CODE - ERROR_QUIET - OUTPUT_STRIP_TRAILING_WHITESPACE) - - if(NOT PYTHON_ERROR_CODE EQUAL 0) - set(${SUCCESS} FALSE PARENT_SCOPE) - set(${OUT} "" PARENT_SCOPE) - else() - set(${SUCCESS} TRUE PARENT_SCOPE) - set(${OUT} ${PYTHON_OUT} PARENT_SCOPE) - endif() -endfunction() - -# Run `EXPR` in python after importing `PKG`. Use the result of this to extend -# `CMAKE_PREFIX_PATH` so the torch cmake configuration can be imported. -macro (append_cmake_prefix_path PKG EXPR) - run_python(_PREFIX_PATH - "import ${PKG}; print(${EXPR})" "Failed to locate ${PKG} path") - list(APPEND CMAKE_PREFIX_PATH ${_PREFIX_PATH}) -endmacro() - -# -# Add a target named `hipify${NAME}` that runs the hipify preprocessor on a set -# of CUDA source files. The names of the corresponding "hipified" sources are -# stored in `OUT_SRCS`. -# -function (hipify_sources_target OUT_SRCS NAME ORIG_SRCS) - # - # Split into C++ and non-C++ (i.e. CUDA) sources. - # - set(NODUP_SRCS ${ORIG_SRCS}) - list(REMOVE_DUPLICATES NODUP_SRCS) - set(SRCS ${NODUP_SRCS}) - set(CXX_SRCS ${NODUP_SRCS}) - list(FILTER SRCS INCLUDE REGEX "\.cu$") - list(FILTER CXX_SRCS EXCLUDE REGEX "\.cu$") - - # - # Generate ROCm/HIP source file names from CUDA file names. - # Since HIP files are generated code, they will appear in the build area - # `CMAKE_CURRENT_BINARY_DIR` directory rather than the original csrc dir. - # - set(HIP_SRCS) - foreach (SRC ${SRCS}) - get_source_file_property(include_dirs "${SRC}" INCLUDE_DIRECTORIES) - get_source_file_property(compile_options "${SRC}" COMPILE_OPTIONS) - string(REGEX REPLACE "\.cu$" "\.hip" SRC ${SRC}) - string(REGEX REPLACE "cuda" "hip" SRC ${SRC}) - - if(include_dirs) - # Copy over include directories from the original CUDA file. - set_source_files_properties( - ${SRC} - PROPERTIES INCLUDE_DIRECTORIES "${include_dirs}") - endif() - - if(compile_options) - set_source_files_properties( - ${SRC} - PROPERTIES COMPILE_OPTIONS "${compile_options}") - endif() - - list(APPEND HIP_SRCS "${CMAKE_CURRENT_BINARY_DIR}/${SRC}") - endforeach() - - add_custom_target( - hipify${NAME} - COMMAND "${Python3_EXECUTABLE}" ${CMAKE_SOURCE_DIR}/cmake/hipify.py -p ${CMAKE_SOURCE_DIR} -o ${CMAKE_CURRENT_BINARY_DIR} ${SRCS} - DEPENDS ${CMAKE_SOURCE_DIR}/cmake/hipify.py ${SRCS} - BYPRODUCTS ${HIP_SRCS} - COMMENT "Running hipify on ${NAME} extension source files.") - - # Swap out original extension sources with hipified sources. - list(APPEND HIP_SRCS ${CXX_SRCS}) - set(${OUT_SRCS} ${HIP_SRCS} PARENT_SCOPE) -endfunction() - -# -# Get additional GPU compiler flags from torch. -# -function (get_torch_gpu_compiler_flags OUT_GPU_FLAGS GPU_LANG) - if (${GPU_LANG} STREQUAL "CUDA") - # - # Get common NVCC flags from torch. - # - run_python(GPU_FLAGS - "from torch.utils.cpp_extension import COMMON_NVCC_FLAGS; print(';'.join(COMMON_NVCC_FLAGS))" - "Failed to determine torch nvcc compiler flags") - - if (CUDA_VERSION VERSION_GREATER_EQUAL 11.8) - list(APPEND GPU_FLAGS "-DENABLE_FP8") - list(REMOVE_ITEM GPU_FLAGS - "-D__CUDA_NO_HALF_OPERATORS__" - "-D__CUDA_NO_HALF_CONVERSIONS__" - "-D__CUDA_NO_BFLOAT16_CONVERSIONS__" - "-D__CUDA_NO_HALF2_OPERATORS__") - endif() - - elseif(${GPU_LANG} STREQUAL "HIP") - # - # Get common HIP/HIPCC flags from torch. - # - run_python(GPU_FLAGS - "import torch.utils.cpp_extension as t; print(';'.join(t.COMMON_HIP_FLAGS + t.COMMON_HIPCC_FLAGS))" - "Failed to determine torch nvcc compiler flags") - - list(APPEND GPU_FLAGS - "-DUSE_ROCM" - "-DENABLE_FP8" - "-U__HIP_NO_HALF_CONVERSIONS__" - "-U__HIP_NO_HALF_OPERATORS__" - "-fno-gpu-rdc") - - endif() - set(${OUT_GPU_FLAGS} ${GPU_FLAGS} PARENT_SCOPE) -endfunction() - -# Macro for converting a `gencode` version number to a cmake version number. -macro(string_to_ver OUT_VER IN_STR) - string(REGEX REPLACE "\([0-9]+\)\([0-9]\)" "\\1.\\2" ${OUT_VER} ${IN_STR}) -endmacro() - -# -# Clear all `-gencode` flags from `CMAKE_CUDA_FLAGS`. -# -# Example: -# CMAKE_CUDA_FLAGS="-Wall -gencode arch=compute_70,code=sm_70 -gencode arch=compute_75,code=sm_75" -# clear_gencode_flags() -# CMAKE_CUDA_FLAGS="-Wall" -# -macro(clear_gencode_flags) - # Remove all `-gencode` flags from `CMAKE_CUDA_FLAGS` since they will be modified - # and passed back via the `CUDA_ARCHITECTURES` property. - string(REGEX REPLACE "-gencode arch=[^ ]+ *" "" CMAKE_CUDA_FLAGS - ${CMAKE_CUDA_FLAGS}) -endmacro() - -# -# Extract unique CUDA architectures from a list of compute capabilities codes in -# the form `[]`, convert them to the form sort -# `.`, dedupes them and then sorts them in ascending order and -# stores them in `OUT_ARCHES`. -# -# Example: -# CUDA_ARCH_FLAGS="-gencode arch=compute_75,code=sm_75;...;-gencode arch=compute_90a,code=sm_90a" -# extract_unique_cuda_archs_ascending(OUT_ARCHES CUDA_ARCH_FLAGS) -# OUT_ARCHES="7.5;...;9.0" -function(extract_unique_cuda_archs_ascending OUT_ARCHES CUDA_ARCH_FLAGS) - set(_CUDA_ARCHES) - foreach(_ARCH ${CUDA_ARCH_FLAGS}) - string(REGEX MATCH "arch=compute_\([0-9]+a?\)" _COMPUTE ${_ARCH}) - if (_COMPUTE) - set(_COMPUTE ${CMAKE_MATCH_1}) - endif() - - string_to_ver(_COMPUTE_VER ${_COMPUTE}) - list(APPEND _CUDA_ARCHES ${_COMPUTE_VER}) - endforeach() - - list(REMOVE_DUPLICATES _CUDA_ARCHES) - list(SORT _CUDA_ARCHES COMPARE NATURAL ORDER ASCENDING) - set(${OUT_ARCHES} ${_CUDA_ARCHES} PARENT_SCOPE) -endfunction() - -# -# For a specific file set the `-gencode` flag in compile options conditionally -# for the CUDA language. -# -# Example: -# set_gencode_flag_for_srcs( -# SRCS "foo.cu" -# ARCH "compute_75" -# CODE "sm_75") -# adds: "-gencode arch=compute_75,code=sm_75" to the compile options for -# `foo.cu` (only for the CUDA language). -# -macro(set_gencode_flag_for_srcs) - set(options) - set(oneValueArgs ARCH CODE) - set(multiValueArgs SRCS) - cmake_parse_arguments(arg "${options}" "${oneValueArgs}" - "${multiValueArgs}" ${ARGN} ) - set(_FLAG -gencode arch=${arg_ARCH},code=${arg_CODE}) - set_property( - SOURCE ${arg_SRCS} - APPEND PROPERTY - COMPILE_OPTIONS "$<$:${_FLAG}>" - ) - - message(DEBUG "Setting gencode flag for ${arg_SRCS}: ${_FLAG}") -endmacro(set_gencode_flag_for_srcs) - -# -# For a list of source files set the `-gencode` flags in the files specific -# compile options (specifically for the CUDA language). -# -# arguments are: -# SRCS: list of source files -# CUDA_ARCHS: list of CUDA architectures in the form `.[letter]` -# BUILD_PTX_FOR_ARCH: if set to true, then the PTX code will be built -# for architecture `BUILD_PTX_FOR_ARCH` if there is a CUDA_ARCH in CUDA_ARCHS -# that is larger than BUILD_PTX_FOR_ARCH. -# -macro(set_gencode_flags_for_srcs) - set(options) - set(oneValueArgs BUILD_PTX_FOR_ARCH) - set(multiValueArgs SRCS CUDA_ARCHS) - cmake_parse_arguments(arg "${options}" "${oneValueArgs}" - "${multiValueArgs}" ${ARGN} ) - - foreach(_ARCH ${arg_CUDA_ARCHS}) - # handle +PTX suffix: generate both sm and ptx codes if requested - string(FIND "${_ARCH}" "+PTX" _HAS_PTX) - if(NOT _HAS_PTX EQUAL -1) - string(REPLACE "+PTX" "" _BASE_ARCH "${_ARCH}") - string(REPLACE "." "" _STRIPPED_ARCH "${_BASE_ARCH}") - set_gencode_flag_for_srcs( - SRCS ${arg_SRCS} - ARCH "compute_${_STRIPPED_ARCH}" - CODE "sm_${_STRIPPED_ARCH}") - set_gencode_flag_for_srcs( - SRCS ${arg_SRCS} - ARCH "compute_${_STRIPPED_ARCH}" - CODE "compute_${_STRIPPED_ARCH}") - else() - string(REPLACE "." "" _STRIPPED_ARCH "${_ARCH}") - set_gencode_flag_for_srcs( - SRCS ${arg_SRCS} - ARCH "compute_${_STRIPPED_ARCH}" - CODE "sm_${_STRIPPED_ARCH}") - endif() - endforeach() - - if (${arg_BUILD_PTX_FOR_ARCH}) - list(SORT arg_CUDA_ARCHS COMPARE NATURAL ORDER ASCENDING) - list(GET arg_CUDA_ARCHS -1 _HIGHEST_ARCH) - if (_HIGHEST_ARCH VERSION_GREATER_EQUAL ${arg_BUILD_PTX_FOR_ARCH}) - string(REPLACE "." "" _PTX_ARCH "${arg_BUILD_PTX_FOR_ARCH}") - set_gencode_flag_for_srcs( - SRCS ${arg_SRCS} - ARCH "compute_${_PTX_ARCH}" - CODE "compute_${_PTX_ARCH}") - endif() - endif() -endmacro() - -# -# For the given `SRC_CUDA_ARCHS` list of gencode versions in the form -# `.[letter]` compute the "loose intersection" with the -# `TGT_CUDA_ARCHS` list of gencodes. We also support the `+PTX` suffix in -# `SRC_CUDA_ARCHS` which indicates that the PTX code should be built when there -# is a CUDA_ARCH in `TGT_CUDA_ARCHS` that is equal to or larger than the -# architecture in `SRC_CUDA_ARCHS`. -# The loose intersection is defined as: -# { max{ x \in tgt | x <= y } | y \in src, { x \in tgt | x <= y } != {} } -# where `<=` is the version comparison operator. -# In other words, for each version in `TGT_CUDA_ARCHS` find the highest version -# in `SRC_CUDA_ARCHS` that is less or equal to the version in `TGT_CUDA_ARCHS`. -# We have special handling for x.0a, if x.0a is in `SRC_CUDA_ARCHS` and x.0 is -# in `TGT_CUDA_ARCHS` then we should remove x.0a from `SRC_CUDA_ARCHS` and add -# x.0a to the result (and remove x.0 from TGT_CUDA_ARCHS). -# The result is stored in `OUT_CUDA_ARCHS`. -# -# Example: -# SRC_CUDA_ARCHS="7.5;8.0;8.6;9.0;9.0a" -# TGT_CUDA_ARCHS="8.0;8.9;9.0" -# cuda_archs_loose_intersection(OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_ARCHS) -# OUT_CUDA_ARCHS="8.0;8.6;9.0;9.0a" -# -# Example With PTX: -# SRC_CUDA_ARCHS="8.0+PTX" -# TGT_CUDA_ARCHS="9.0" -# cuda_archs_loose_intersection(OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_ARCHS) -# OUT_CUDA_ARCHS="8.0+PTX" -# -function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_ARCHS) - set(_SRC_CUDA_ARCHS "${SRC_CUDA_ARCHS}") - set(_TGT_CUDA_ARCHS ${TGT_CUDA_ARCHS}) - - # handle +PTX suffix: separate base arch for matching, record PTX requests - set(_PTX_ARCHS) - foreach(_arch ${_SRC_CUDA_ARCHS}) - if(_arch MATCHES "\\+PTX$") - string(REPLACE "+PTX" "" _base "${_arch}") - list(APPEND _PTX_ARCHS "${_base}") - list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}") - list(APPEND _SRC_CUDA_ARCHS "${_base}") - endif() - endforeach() - list(REMOVE_DUPLICATES _PTX_ARCHS) - list(REMOVE_DUPLICATES _SRC_CUDA_ARCHS) - - # If x.0a or x.0f is in SRC_CUDA_ARCHS and x.0 is in CUDA_ARCHS then we should - # remove x.0a or x.0f from SRC_CUDA_ARCHS and add x.0a or x.0f to _CUDA_ARCHS - set(_CUDA_ARCHS) - foreach(_arch ${_SRC_CUDA_ARCHS}) - if(_arch MATCHES "[af]$") - list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}") - string(REGEX REPLACE "[af]$" "" _base "${_arch}") - if ("${_base}" IN_LIST TGT_CUDA_ARCHS) - list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}") - list(APPEND _CUDA_ARCHS "${_arch}") - endif() - endif() - endforeach() - - list(SORT _SRC_CUDA_ARCHS COMPARE NATURAL ORDER ASCENDING) - - # for each ARCH in TGT_CUDA_ARCHS find the highest arch in SRC_CUDA_ARCHS that - # is less or equal to ARCH (but has the same major version since SASS binary - # compatibility is only forward compatible within the same major version). - foreach(_ARCH ${_TGT_CUDA_ARCHS}) - set(_TMP_ARCH) - # Extract the major version of the target arch - string(REGEX REPLACE "^([0-9]+)\\..*$" "\\1" TGT_ARCH_MAJOR "${_ARCH}") - foreach(_SRC_ARCH ${_SRC_CUDA_ARCHS}) - # Extract the major version of the source arch - string(REGEX REPLACE "^([0-9]+)\\..*$" "\\1" SRC_ARCH_MAJOR "${_SRC_ARCH}") - # Check version-less-or-equal, and allow PTX arches to match across majors - if (_SRC_ARCH VERSION_LESS_EQUAL _ARCH) - if (_SRC_ARCH IN_LIST _PTX_ARCHS OR SRC_ARCH_MAJOR STREQUAL TGT_ARCH_MAJOR) - set(_TMP_ARCH "${_SRC_ARCH}") - endif() - else() - # If we hit a version greater than the target, we can break - break() - endif() - endforeach() - - # If we found a matching _TMP_ARCH, append it to _CUDA_ARCHS - if (_TMP_ARCH) - list(APPEND _CUDA_ARCHS "${_TMP_ARCH}") - endif() - endforeach() - - list(REMOVE_DUPLICATES _CUDA_ARCHS) - - # reapply +PTX suffix to architectures that requested PTX - set(_FINAL_ARCHS) - foreach(_arch ${_CUDA_ARCHS}) - if(_arch IN_LIST _PTX_ARCHS) - list(APPEND _FINAL_ARCHS "${_arch}+PTX") - else() - list(APPEND _FINAL_ARCHS "${_arch}") - endif() - endforeach() - set(_CUDA_ARCHS ${_FINAL_ARCHS}) - - list(SORT _CUDA_ARCHS COMPARE NATURAL ORDER ASCENDING) - - set(${OUT_CUDA_ARCHS} ${_CUDA_ARCHS} PARENT_SCOPE) -endfunction() - -# -# For the given `SRC_ROCM_ARCHS` list of architecture versions in the form -# `` compute the "loose intersection" with the `TGT_ROCM_ARCHS` list. -# The loose intersection is defined as: -# { max{ x \in tgt | x <= y } | y \in src, { x \in tgt | x <= y } != {} } -# where `<=` is the version comparison operator. -# In other words, for each version in `TGT_ROCM_ARCHS` find the highest version -# in `SRC_ROCM_ARCHS` that is less or equal to the version in `TGT_ROCM_ARCHS`. -# The result is stored in `OUT_ROCM_ARCHS`. -# -# Example: -# SRC_ROCM_ARCHS="gfx900;gfx906;gfx908;gfx90a" -# TGT_ROCM_ARCHS="gfx906;gfx908;gfx1030" -# hip_archs_loose_intersection(OUT_ROCM_ARCHS SRC_ROCM_ARCHS TGT_ROCM_ARCHS) -# OUT_ROCM_ARCHS="gfx906;gfx908" -# -function(hip_archs_loose_intersection OUT_ROCM_ARCHS SRC_ROCM_ARCHS TGT_ROCM_ARCHS) - list(REMOVE_DUPLICATES SRC_ROCM_ARCHS) - - # ROCm architectures are typically in format gfxNNN or gfxNNNx where N is a digit - # and x is a letter. We can sort them by string comparison which works for this format. - list(SORT SRC_ROCM_ARCHS COMPARE STRING ORDER ASCENDING) - - set(_ROCM_ARCHS) - - # Find the intersection of supported architectures - foreach(_SRC_ARCH ${SRC_ROCM_ARCHS}) - if(_SRC_ARCH IN_LIST TGT_ROCM_ARCHS) - list(APPEND _ROCM_ARCHS ${_SRC_ARCH}) - endif() - endforeach() - - list(REMOVE_DUPLICATES _ROCM_ARCHS) - set(${OUT_ROCM_ARCHS} ${_ROCM_ARCHS} PARENT_SCOPE) -endfunction() - -function(cuda_remove_ptx_suffixes OUT_CUDA_ARCHS CUDA_ARCHS) - set(_CUDA_ARCHS "${CUDA_ARCHS}") - - # handle +PTX suffix: separate base arch for matching, record PTX requests - foreach(_arch ${CUDA_ARCHS}) - if(_arch MATCHES "\\+PTX$") - string(REPLACE "+PTX" "" _base "${_arch}") - list(REMOVE_ITEM _CUDA_ARCHS "${_arch}") - list(APPEND _CUDA_ARCHS "${_base}") - endif() - endforeach() - - list(REMOVE_DUPLICATES _CUDA_ARCHS) - list(SORT _CUDA_ARCHS COMPARE NATURAL ORDER ASCENDING) - - set(${OUT_CUDA_ARCHS} ${_CUDA_ARCHS} PARENT_SCOPE) -endfunction() - -# -# Define a target named `GPU_MOD_NAME` for a single extension. The -# arguments are: -# -# DESTINATION - Module destination directory. -# LANGUAGE - The GPU language for this module, e.g CUDA, HIP, -# etc. -# SOURCES - List of source files relative to CMakeLists.txt -# directory. -# -# Optional arguments: -# -# ARCHITECTURES - A list of target GPU architectures in cmake -# format. -# Refer `CMAKE_CUDA_ARCHITECTURES` documentation -# and `CMAKE_HIP_ARCHITECTURES` for more info. -# ARCHITECTURES will use cmake's defaults if -# not provided. -# COMPILE_FLAGS - Extra compiler flags passed to NVCC/hip. -# INCLUDE_DIRECTORIES - Extra include directories. -# LIBRARIES - Extra link libraries. -# WITH_SOABI - Generate library with python SOABI suffix name. -# USE_SABI - Use python stable api -# -# Note: optimization level/debug info is set via cmake build type. -# -function (define_gpu_extension_target GPU_MOD_NAME) - cmake_parse_arguments(PARSE_ARGV 1 - GPU - "WITH_SOABI" - "DESTINATION;LANGUAGE;USE_SABI" - "SOURCES;COMPILE_FLAGS;INCLUDE_DIRECTORIES;LIBRARIES") - - # Add hipify preprocessing step when building with HIP/ROCm. - if (GPU_LANGUAGE STREQUAL "HIP") - hipify_sources_target(GPU_SOURCES ${GPU_MOD_NAME} "${GPU_SOURCES}") - endif() - - if (GPU_WITH_SOABI) - set(GPU_WITH_SOABI WITH_SOABI) - else() - set(GPU_WITH_SOABI) - endif() - - if (GPU_USE_SABI) - Python3_add_library(${GPU_MOD_NAME} MODULE USE_SABI ${GPU_USE_SABI} ${GPU_WITH_SOABI} "${GPU_SOURCES}") - else() - Python3_add_library(${GPU_MOD_NAME} MODULE ${GPU_WITH_SOABI} "${GPU_SOURCES}") - endif() - - if (GPU_LANGUAGE STREQUAL "HIP") - # Make this target dependent on the hipify preprocessor step. - add_dependencies(${GPU_MOD_NAME} hipify${GPU_MOD_NAME}) - - # Clear target architectures, we are passing arch flags per source file. - set_property(TARGET ${GPU_MOD_NAME} PROPERTY HIP_ARCHITECTURES off) - endif() - - if (TORCH_VERSION VERSION_LESS 2.12.0) - set_property(TARGET ${GPU_MOD_NAME} PROPERTY CXX_STANDARD 17) - else() - set_property(TARGET ${GPU_MOD_NAME} PROPERTY CXX_STANDARD 20) - endif() - - target_compile_options(${GPU_MOD_NAME} PRIVATE - $<$:${GPU_COMPILE_FLAGS}>) - - target_compile_definitions(${GPU_MOD_NAME} PRIVATE - "-DTORCH_EXTENSION_NAME=${GPU_MOD_NAME}") - - target_include_directories(${GPU_MOD_NAME} PRIVATE csrc - ${GPU_INCLUDE_DIRECTORIES}) - - target_link_libraries(${GPU_MOD_NAME} PRIVATE torch ${GPU_LIBRARIES}) - - # Don't use `TORCH_LIBRARIES` for CUDA since it pulls in a bunch of - # dependencies that are not necessary and may not be installed. - if (GPU_LANGUAGE STREQUAL "CUDA") - target_link_libraries(${GPU_MOD_NAME} PRIVATE CUDA::cudart) - else() - target_link_libraries(${GPU_MOD_NAME} PRIVATE ${TORCH_LIBRARIES}) - endif() - - install(TARGETS ${GPU_MOD_NAME} LIBRARY DESTINATION ${GPU_DESTINATION} COMPONENT ${GPU_MOD_NAME}) -endfunction() - -# Map a GPU language to its backend name. -# -# Arguments: -# OUT_BACKEND - Output variable name for the backend string -# GPU_LANG - The GPU language (CPU, CUDA, HIP, METAL, SYCL) -# -function(gpu_lang_to_backend OUT_BACKEND GPU_LANG) - if (${GPU_LANG} STREQUAL "CPU") - set(_BACKEND "cpu") - elseif (${GPU_LANG} STREQUAL "CUDA") - set(_BACKEND "cuda") - elseif (${GPU_LANG} STREQUAL "HIP") - set(_BACKEND "rocm") - elseif (${GPU_LANG} STREQUAL "METAL") - set(_BACKEND "metal") - elseif (${GPU_LANG} STREQUAL "SYCL") - set(_BACKEND "xpu") - else() - message(FATAL_ERROR "Unsupported GPU_LANG: ${GPU_LANG}") - endif() - - set(${OUT_BACKEND} "${_BACKEND}" PARENT_SCOPE) -endfunction() diff --git a/natten/compat.py b/natten/compat.py deleted file mode 100644 index a9b2672c..00000000 --- a/natten/compat.py +++ /dev/null @@ -1,26 +0,0 @@ -import ctypes -import importlib.util -import sys -from pathlib import Path -from types import ModuleType - - -def _import_from_path(file_path: Path) -> ModuleType: - # We cannot use the module name as-is, after adding it to `sys.modules`, - # it would also be used for other imports. So, we make a module name that - # depends on the path for it to be unique using the hex-encoded hash of - # the path. - path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) - module_name = path_hash - spec = importlib.util.spec_from_file_location(module_name, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[module_name] = module - spec.loader.exec_module(module) # type: ignore - return module - - -globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/natten/metadata-cann.json b/natten/metadata-cann.json deleted file mode 100644 index 3123bb7b..00000000 --- a/natten/metadata-cann.json +++ /dev/null @@ -1,22 +0,0 @@ -{ - "name": "natten", - "id": "_natten_cann_28fa1dd", - "version": 1, - "license": "MIT", - "upstream": "https://github.com/SHI-Labs/NATTEN", - "python-depends": [], - "backend": { - "type": "cann" - }, - "provenance": { - "kernel-builder": { - "version": "0.17.0-dev0", - "sha": "81580bb92577f2f7228661ca2a221fc052375709", - "dirty": true - }, - "kernel": { - "sha": "39160cf398fdf251faf03ede7fa8d56dfa29701e", - "dirty": false - } - } -} \ No newline at end of file diff --git a/natten/metadata-cpu.json b/natten/metadata-cpu.json deleted file mode 100644 index 68482e01..00000000 --- a/natten/metadata-cpu.json +++ /dev/null @@ -1,22 +0,0 @@ -{ - "name": "natten", - "id": "_natten_cpu_28fa1dd", - "version": 1, - "license": "MIT", - "upstream": "https://github.com/SHI-Labs/NATTEN", - "python-depends": [], - "backend": { - "type": "cpu" - }, - "provenance": { - "kernel-builder": { - "version": "0.17.0-dev0", - "sha": "81580bb92577f2f7228661ca2a221fc052375709", - "dirty": true - }, - "kernel": { - "sha": "39160cf398fdf251faf03ede7fa8d56dfa29701e", - "dirty": false - } - } -} \ No newline at end of file diff --git a/natten/metadata-cuda.json b/natten/metadata-cuda.json deleted file mode 100644 index 5771d35d..00000000 --- a/natten/metadata-cuda.json +++ /dev/null @@ -1,22 +0,0 @@ -{ - "name": "natten", - "id": "_natten_cuda_28fa1dd", - "version": 1, - "license": "MIT", - "upstream": "https://github.com/SHI-Labs/NATTEN", - "python-depends": [], - "backend": { - "type": "cuda" - }, - "provenance": { - "kernel-builder": { - "version": "0.17.0-dev0", - "sha": "81580bb92577f2f7228661ca2a221fc052375709", - "dirty": true - }, - "kernel": { - "sha": "39160cf398fdf251faf03ede7fa8d56dfa29701e", - "dirty": false - } - } -} \ No newline at end of file diff --git a/natten/metadata-metal.json b/natten/metadata-metal.json deleted file mode 100644 index 7fd48375..00000000 --- a/natten/metadata-metal.json +++ /dev/null @@ -1,22 +0,0 @@ -{ - "name": "natten", - "id": "_natten_metal_28fa1dd", - "version": 1, - "license": "MIT", - "upstream": "https://github.com/SHI-Labs/NATTEN", - "python-depends": [], - "backend": { - "type": "metal" - }, - "provenance": { - "kernel-builder": { - "version": "0.17.0-dev0", - "sha": "81580bb92577f2f7228661ca2a221fc052375709", - "dirty": true - }, - "kernel": { - "sha": "39160cf398fdf251faf03ede7fa8d56dfa29701e", - "dirty": false - } - } -} \ No newline at end of file diff --git a/natten/metadata-neuron.json b/natten/metadata-neuron.json deleted file mode 100644 index 7bed17df..00000000 --- a/natten/metadata-neuron.json +++ /dev/null @@ -1,22 +0,0 @@ -{ - "name": "natten", - "id": "_natten_neuron_28fa1dd", - "version": 1, - "license": "MIT", - "upstream": "https://github.com/SHI-Labs/NATTEN", - "python-depends": [], - "backend": { - "type": "neuron" - }, - "provenance": { - "kernel-builder": { - "version": "0.17.0-dev0", - "sha": "81580bb92577f2f7228661ca2a221fc052375709", - "dirty": true - }, - "kernel": { - "sha": "39160cf398fdf251faf03ede7fa8d56dfa29701e", - "dirty": false - } - } -} \ No newline at end of file diff --git a/natten/metadata-rocm.json b/natten/metadata-rocm.json deleted file mode 100644 index 990b7bf9..00000000 --- a/natten/metadata-rocm.json +++ /dev/null @@ -1,22 +0,0 @@ -{ - "name": "natten", - "id": "_natten_rocm_28fa1dd", - "version": 1, - "license": "MIT", - "upstream": "https://github.com/SHI-Labs/NATTEN", - "python-depends": [], - "backend": { - "type": "rocm" - }, - "provenance": { - "kernel-builder": { - "version": "0.17.0-dev0", - "sha": "81580bb92577f2f7228661ca2a221fc052375709", - "dirty": true - }, - "kernel": { - "sha": "39160cf398fdf251faf03ede7fa8d56dfa29701e", - "dirty": false - } - } -} \ No newline at end of file diff --git a/natten/metadata-xpu.json b/natten/metadata-xpu.json deleted file mode 100644 index 8409363a..00000000 --- a/natten/metadata-xpu.json +++ /dev/null @@ -1,22 +0,0 @@ -{ - "name": "natten", - "id": "_natten_xpu_28fa1dd", - "version": 1, - "license": "MIT", - "upstream": "https://github.com/SHI-Labs/NATTEN", - "python-depends": [], - "backend": { - "type": "xpu" - }, - "provenance": { - "kernel-builder": { - "version": "0.17.0-dev0", - "sha": "81580bb92577f2f7228661ca2a221fc052375709", - "dirty": true - }, - "kernel": { - "sha": "39160cf398fdf251faf03ede7fa8d56dfa29701e", - "dirty": false - } - } -} \ No newline at end of file diff --git a/natten/pyproject.toml b/natten/pyproject.toml deleted file mode 100644 index 3743f4d9..00000000 --- a/natten/pyproject.toml +++ /dev/null @@ -1,23 +0,0 @@ -[project] -name = "natten" -version = "0.1.0" -requires-python = ">=3.9" - -[build-system] -requires = [ - "cmake>=3.26", - "ninja", - "packaging", - "setuptools>=61", - "torch", - "wheel", - -] -build-backend = "setuptools.build_meta" - -[project.optional-dependencies] - -[tool.pytest.ini_options] -markers = [ - "kernels_ci: mark a test as a kernel CI test" -] \ No newline at end of file diff --git a/natten/setup.py b/natten/setup.py deleted file mode 100644 index 92f797d7..00000000 --- a/natten/setup.py +++ /dev/null @@ -1,200 +0,0 @@ -import logging -import os -from shutil import which, move -import subprocess -import sys -from pathlib import Path - -from setuptools import Extension, find_packages, setup -from setuptools.command.build import build -from setuptools.command.build_ext import build_ext - -logger = logging.getLogger(__name__) - - -def get_backend() -> str: - """Detect the backend by inspecting torch.""" - import torch - - if torch.version.cuda is not None: - return "cuda" - elif torch.version.hip is not None: - return "rocm" - elif torch.backends.mps.is_available(): - return "metal" - elif hasattr(torch.version, "xpu") and torch.version.xpu is not None: - return "xpu" - else: - return "cpu" - - -def is_sccache_available() -> bool: - return which("sccache") is not None - - -def is_ccache_available() -> bool: - return which("ccache") is not None - - -def is_ninja_available() -> bool: - return which("ninja") is not None - - -def _make_cmake_args(cfg: str) -> tuple[list[str], list[str]]: - """Build CMake and build arguments from the current environment.""" - cmake_generator = os.environ.get("CMAKE_GENERATOR", "") - - cmake_args = [ - f"-DPython3_EXECUTABLE={sys.executable}", - f"-DCMAKE_BUILD_TYPE={cfg}", # not used on MSVC, but no harm - ] - build_args: list[str] = [] - - if "CMAKE_ARGS" in os.environ: - cmake_args += [item for item in os.environ["CMAKE_ARGS"].split(" ") if item] - - if not cmake_generator or cmake_generator == "Ninja": - try: - import ninja - - ninja_executable_path = Path(ninja.BIN_DIR) / "ninja" - cmake_args += [ - "-GNinja", - f"-DCMAKE_MAKE_PROGRAM:FILEPATH={ninja_executable_path}", - ] - except ImportError: - pass - - if is_sccache_available(): - cmake_args += [ - "-DCMAKE_C_COMPILER_LAUNCHER=sccache", - "-DCMAKE_CXX_COMPILER_LAUNCHER=sccache", - "-DCMAKE_CUDA_COMPILER_LAUNCHER=sccache", - "-DCMAKE_HIP_COMPILER_LAUNCHER=sccache", - "-DCMAKE_OBJC_COMPILER_LAUNCHER=sccache", - "-DCMAKE_OBJCXX_COMPILER_LAUNCHER=sccache", - ] - elif is_ccache_available(): - cmake_args += [ - "-DCMAKE_C_COMPILER_LAUNCHER=ccache", - "-DCMAKE_CXX_COMPILER_LAUNCHER=ccache", - "-DCMAKE_CUDA_COMPILER_LAUNCHER=ccache", - "-DCMAKE_HIP_COMPILER_LAUNCHER=ccache", - "-DCMAKE_OBJC_COMPILER_LAUNCHER=ccache", - "-DCMAKE_OBJCXX_COMPILER_LAUNCHER=ccache", - ] - - num_jobs = os.getenv("MAX_JOBS", None) - if num_jobs is not None: - num_jobs = int(num_jobs) - logger.info("Using MAX_JOBS=%d as the number of jobs.", num_jobs) - else: - try: - # os.sched_getaffinity() isn't universally available, so fall - # back to os.cpu_count() if we get an error here. - num_jobs = len(os.sched_getaffinity(0)) - except AttributeError: - num_jobs = os.cpu_count() - - nvcc_threads = os.getenv("NVCC_THREADS", None) - if nvcc_threads is not None: - nvcc_threads = int(nvcc_threads) - logger.info( - "Using NVCC_THREADS=%d as the number of nvcc threads.", nvcc_threads - ) - num_jobs = max(1, num_jobs // nvcc_threads) - cmake_args += ["-DNVCC_THREADS={}".format(nvcc_threads)] - - build_args += [f"-j{num_jobs}"] - if sys.platform == "win32": - build_args += ["--config", cfg] - - return cmake_args, build_args - - -class CMakeExtension(Extension): - def __init__(self, name: str, sourcedir: str = "") -> None: - super().__init__(name, sources=[], py_limited_api=True) - self.sourcedir = os.fspath(Path(sourcedir).resolve()) - - -class CMakeBuild(build_ext): - def build_extension(self, ext: CMakeExtension) -> None: - ext_fullpath = Path.cwd() / self.get_ext_fullpath(ext.name) - extdir = ext_fullpath.parent.resolve() - - debug = int(os.environ.get("DEBUG", 0)) if self.debug is None else self.debug - cfg = "Debug" if debug else "Release" - - cmake_args, build_args = _make_cmake_args(cfg) - cmake_args = [f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={extdir}{os.sep}"] + cmake_args - - build_temp = Path(self.build_temp) / ext.name - if not build_temp.exists(): - build_temp.mkdir(parents=True) - - subprocess.run( - ["cmake", "-S", ext.sourcedir, "-B", str(build_temp), *cmake_args], - cwd=build_temp, - check=True, - ) - subprocess.run( - ["cmake", "--build", str(build_temp), *build_args], cwd=build_temp, check=True - ) - - if sys.platform == "win32": - # Move the dylib one folder up for discovery. - for filename in os.listdir(extdir / cfg): - move(extdir / cfg / filename, extdir / filename) - - -class BuildKernel(build): - """Custom command to build and locally install the kernel.""" - - description = "Build the kernel and install via the local_install CMake target" - user_options = [] - - def initialize_options(self) -> None: - super().initialize_options() - - def finalize_options(self) -> None: - super().finalize_options() - - def run(self) -> None: - project_root = Path(__file__).parent - - debug = int(os.environ.get("DEBUG", 0)) - cfg = "Debug" if debug else "Release" - - cmake_args, build_args = _make_cmake_args(cfg) - - build_temp = project_root / "_cmake_build" - build_temp.mkdir(parents=True, exist_ok=True) - - subprocess.run( - ["cmake", "-S", str(project_root), "-B", str(build_temp), *cmake_args], - cwd=project_root, - check=True, - ) - subprocess.run( - ["cmake", "--build", str(build_temp), "--target", "local_install", *build_args], - cwd=project_root, - check=True, - ) - - -backend = get_backend() -ops_name = f"_natten_{backend}_28fa1dd" - -setup( - name="natten", - # The version is just a stub, it's not used by the final build artefact. - version="0.1.0", - ext_modules=[CMakeExtension(f"natten.{ops_name}")], - cmdclass={"build_ext": CMakeBuild, "build_kernel": BuildKernel}, - packages=find_packages(where="torch-ext", include=["natten*"]), - package_dir={"": "torch-ext"}, - zip_safe=False, - install_requires=["torch"], - python_requires=">=3.9", -) \ No newline at end of file diff --git a/natten/torch-ext/registration.h b/natten/torch-ext/registration.h deleted file mode 100644 index 096da4be..00000000 --- a/natten/torch-ext/registration.h +++ /dev/null @@ -1,39 +0,0 @@ -// Registration macros from vLLM: -// https://github.com/vllm-project/vllm/blob/main/csrc/core/registration.h - -#pragma once - -#include - -#define _CONCAT(A, B) A##B -#define CONCAT(A, B) _CONCAT(A, B) - -#define _STRINGIFY(A) #A -#define STRINGIFY(A) _STRINGIFY(A) - -// A version of the TORCH_LIBRARY macro that expands the NAME, i.e. so NAME -// could be a macro instead of a literal token. -#define TORCH_LIBRARY_EXPAND(NAME, MODULE) TORCH_LIBRARY(NAME, MODULE) - -// A version of the STABLE_TORCH_LIBRARY macro that expands the NAME, i.e. so NAME -// could be a macro instead of a literal token. -#define STABLE_TORCH_LIBRARY_EXPAND(NAME, MODULE) STABLE_TORCH_LIBRARY(NAME, MODULE) - -// A version of the TORCH_LIBRARY_IMPL macro that expands the NAME, i.e. so NAME -// could be a macro instead of a literal token. -#define TORCH_LIBRARY_IMPL_EXPAND(NAME, DEVICE, MODULE) \ - TORCH_LIBRARY_IMPL(NAME, DEVICE, MODULE) - -// A version of the STABLE_TORCH_LIBRARY_IMPL macro that expands the NAME, i.e. so NAME -// could be a macro instead of a literal token. -#define STABLE_TORCH_LIBRARY_IMPL_EXPAND(NAME, DEVICE, MODULE) \ - STABLE_TORCH_LIBRARY_IMPL(NAME, DEVICE, MODULE) - -// REGISTER_EXTENSION allows the shared library to be loaded and initialized -// via python's import statement. -#define REGISTER_EXTENSION(NAME) \ - PyMODINIT_FUNC CONCAT(PyInit_, NAME)() { \ - static struct PyModuleDef module = {PyModuleDef_HEAD_INIT, \ - STRINGIFY(NAME), nullptr, 0, nullptr}; \ - return PyModule_Create(&module); \ - } From f56f26d3018b4d000f40b04551e0ab9b7586988f Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Fri, 17 Jul 2026 17:55:50 +0530 Subject: [PATCH 7/8] natten: bridge compat-shim submodule imports with a meta-path finder kernel-builder's compat shim executes the package under a path-derived module name and only copies globals into `natten`, so real submodule imports (from natten.types / natten.utils.testing / ... import X) raised ModuleNotFoundError in the CI test shell. Install a meta-path finder (only when `natten` in sys.modules is our own compat shim) that resolves natten.* to the already-loaded hash-named modules, preserving module identity and avoiding double execution via parent __path__ lookups. Co-Authored-By: Claude Fable 5 --- natten/torch-ext/natten/__init__.py | 60 +++++++++++++++++++++++++++++ 1 file changed, 60 insertions(+) diff --git a/natten/torch-ext/natten/__init__.py b/natten/torch-ext/natten/__init__.py index 9e0a40d5..9fe10d40 100644 --- a/natten/torch-ext/natten/__init__.py +++ b/natten/torch-ext/natten/__init__.py @@ -75,6 +75,66 @@ _sys.modules[__name__ + ".types"] = types +# kernel-builder's compat shim (`natten/__init__.py` inside the build variant +# directory) executes this package under a path-derived module name and copies +# our globals into a `natten` module whose __path__ contains no submodules. +# Attribute access (`natten.functional`) works there, but real submodule +# imports (`from natten.functional import na2d`, `import natten.utils.testing`) +# would either fail with ModuleNotFoundError or — when resolvable through a +# parent package's __path__ — re-execute the module under a second name, +# duplicating module state. Bridge this with a meta-path finder that resolves +# any `natten.*` import to our already-loaded module objects. Only installed +# when `natten` in sys.modules is *our* compat shim, so a real `natten` +# distribution in the same environment is never hijacked. +if __name__ != "natten": + from pathlib import Path as _Path + + _compat = _sys.modules.get("natten") + _is_our_compat = ( + _compat is not None + and getattr(_compat, "__file__", None) is not None + and _Path(_compat.__file__).resolve() + == _Path(__file__).resolve().parent / "natten" / "__init__.py" + ) + + if _is_our_compat: + import importlib as _importlib + from importlib.abc import Loader as _Loader + from importlib.abc import MetaPathFinder as _MetaPathFinder + from importlib.util import spec_from_loader as _spec_from_loader + + _real_root = __name__ + + class _NattenAliasLoader(_Loader): + def __init__(self, module): + self._module = module + self._spec = getattr(module, "__spec__", None) + self._loader = getattr(module, "__loader__", None) + + def create_module(self, spec): + return self._module + + def exec_module(self, module): + # The import machinery stamped the alias spec onto the real + # module in module_from_spec; restore its original identity. + module.__spec__ = self._spec + module.__loader__ = self._loader + + class _NattenAliasFinder(_MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if not fullname.startswith("natten."): + return None + real_name = _real_root + fullname[len("natten") :] + try: + module = _importlib.import_module(real_name) + except ImportError: + return None + return _spec_from_loader(fullname, _NattenAliasLoader(module)) + + # Must precede PathFinder, which would otherwise re-execute + # submodules reachable through a real parent package's __path__. + _sys.meta_path.insert(0, _NattenAliasFinder()) + __all__ = [ "__version__", "NeighborhoodAttention1D", From 73e4fe133875789587406b9ee981842280fcd4cc Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Sat, 18 Jul 2026 08:15:53 +0530 Subject: [PATCH 8/8] natten: adapt Blackwell F8F6F4 MMA atom matching to CUTLASS 4.5 CUTLASS 4.5 turned cute::SM100_MMA_F8F6F4_SS from a tag type (paired with a variadic MMA_Traits specialization) into a class template, same shape as SM100_MMA_F16BF16_SS. The to_tiled_mma_sm100_ts overloads in the fna/fmha Blackwell collectives still pattern-matched the 4.4 traits form, failing on every blackwell TU (first exercised by the cu128+ release builds; PR builds only cover cu126, where the blackwell section is skipped via cuda-minver). Match the templated atom directly, exactly like the adjacent F16BF16 overload, which mirrors what the CUTLASS 4.5 collective builder emits via make_tiled_mma. Co-Authored-By: Claude Fable 5 --- .../fmha_blackwell/collective/fmha_common.hpp | 19 +++++++++++-------- .../fna_blackwell/collective/fna_common.hpp | 19 +++++++++++-------- 2 files changed, 22 insertions(+), 16 deletions(-) diff --git a/natten/csrc/include/natten/cuda/fmha_blackwell/collective/fmha_common.hpp b/natten/csrc/include/natten/cuda/fmha_blackwell/collective/fmha_common.hpp index 6c96030d..76b7e877 100644 --- a/natten/csrc/include/natten/cuda/fmha_blackwell/collective/fmha_common.hpp +++ b/natten/csrc/include/natten/cuda/fmha_blackwell/collective/fmha_common.hpp @@ -95,17 +95,20 @@ template < CUTE_HOST_DEVICE constexpr auto to_tiled_mma_sm100_ts( TiledMMA< MMA_Atom< - MMA_Traits< - SM100_MMA_F8F6F4_SS, + // CUTLASS 4.5 turned SM100_MMA_F8F6F4_SS into a class template + // (upstream NATTEN pins 4.4, where it was a tag type paired with + // a variadic MMA_Traits); match the atom directly, like the + // F16BF16 overload below. + SM100_MMA_F8F6F4_SS< a_type, b_type, c_type, - cute::C, - cute::C, - cute::integral_constant, - cute::integral_constant, - cute::integral_constant, - cute::integral_constant>, + M, + N, + a_major, + b_major, + a_neg, + b_neg>, TAs...>, TMs...>) { return TiledMMA< diff --git a/natten/csrc/include/natten/cuda/fna_blackwell/collective/fna_common.hpp b/natten/csrc/include/natten/cuda/fna_blackwell/collective/fna_common.hpp index 128f3ced..1dbf0658 100644 --- a/natten/csrc/include/natten/cuda/fna_blackwell/collective/fna_common.hpp +++ b/natten/csrc/include/natten/cuda/fna_blackwell/collective/fna_common.hpp @@ -95,17 +95,20 @@ template < CUTE_HOST_DEVICE constexpr auto to_tiled_mma_sm100_ts( TiledMMA< MMA_Atom< - MMA_Traits< - SM100_MMA_F8F6F4_SS, + // CUTLASS 4.5 turned SM100_MMA_F8F6F4_SS into a class template + // (upstream NATTEN pins 4.4, where it was a tag type paired with + // a variadic MMA_Traits); match the atom directly, like the + // F16BF16 overload below. + SM100_MMA_F8F6F4_SS< a_type, b_type, c_type, - cute::C, - cute::C, - cute::integral_constant, - cute::integral_constant, - cute::integral_constant, - cute::integral_constant>, + M, + N, + a_major, + b_major, + a_neg, + b_neg>, TAs...>, TMs...>) { return TiledMMA<